#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Spinner movement equations solved numerically


@author: lbo
"""

import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D  # noqa: F401 unused import
import numpy as np

from collections import namedtuple
import itertools
import math

# MoveParams specifies the vertical angle, its first and second derivative,
# and the angle on the horizontal plane of a spinner.
MoveParams = namedtuple('MoveParams', ['x', 'v', 'a', 'phi'])

Spins = namedtuple('Spins', ['lv', 'ls', 'lorth'])
StdSpins = namedtuple('StdSpins', ['l1', 'l2', 'l3'])

class Calculator:
    _initial = MoveParams(0,0,0,0)
    _lmbda = 0.0
    _nu = 0.0

    _last = MoveParams(0,0,0,0)

    _step = 1e-3

    def __init__(self, initial, nu, lmbda, step=1e-3):
        """Construct a calculator given a set of initial conditions."""
        self._step = step
        assert initial.x != 0, "θ must not be equal to zero in initial conditions"
        # Set initial conditions and calculate initial acceleration.
        self._initial = initial._replace(a=Calculator.acceleration(initial.x, nu, lmbda))
        self._lmbda = lmbda
        self._nu = nu
        self.reset()

    def reset(self):
        self._last = self._initial

    # The calculation methods are static and thus do not change the state
    # of a Calculator instance.
    #
    # ν is not v! It's the greek letter nu. Replace with "nu" after it having
    # caused the second bug.
    def acceleration(θ, ν, λ):
        """The equation for d^2θ/dt^2."""
        first = -(λ - math.cos(θ))*(1 - λ * math.cos(θ)) / (math.sin(θ)**3)
        second = ν * math.sin(θ)
        return first + second

    def next_x(last, step):
        return last.x + last.v * step + last.a * step * step * 1/2

    def next_v(last, a_next, step):
        return last.v + (a_next + last.a)/2 * step

    def incremental_phi(current, last, λ, step):
        """Returns the sum term of phi for the current step."""
        first = (λ - math.cos(current.x))/(math.sin(current.x)**2)
        second = (λ - math.cos(last.x))/(math.sin(last.x)**2)
        return 1/2 * step * (first+second)

    def update(self):
        """Run one iteration."""
        last = self._last
        x_next = Calculator.next_x(last, self._step)
        a_next = Calculator.acceleration(x_next, self._nu, self._lmbda)
        v_next = Calculator.next_v(last, a_next, self._step)

        next = MoveParams(x_next, v_next, a_next, 0)
        # Calculation of phi requires knowing both the current and last value
        # of θ.
        phi_next = last.phi + Calculator.incremental_phi(next, last, self._lmbda, self._step,)
        next = next._replace(phi=phi_next)
        self._last = next
        return last

    def __iter__(self):
        return self

    def __next__(self):
        return self.next()

    def next(self):
        return self.update()

    def spins(self, state=None):
        """Calculate angular momentum in local (e3/e3'/e_orth) base."""
        state = self._last if not state else state
        lv = (self._lmbda - math.cos(state.x))/math.sin(state.x)**2
        ls = (1 - self._lmbda * math.cos(state.x))/math.sin(state.x)**2
        lorth = state.v
        return Spins(lv, ls, lorth)

    def orth_spins(self, state=None):
        """Calculate angular momentum in e1/e2/e3 base. Returns a (l1, l2, l3) tuple."""
        state = self._last if not state else state
        (lv, ls, lorth) = self.spins(state=state)
        # Project local angular momentum vectors onto standard unit vectors e1-e3.
        l1 = ls * math.cos(state.phi) * math.sin(state.x) - lorth * math.sin(state.phi)
        l2 = ls * math.sin(state.phi) * math.sin(state.x) + lorth * math.cos(state.phi)
        l3 = lv + ls * math.cos(state.x)
        return StdSpins(l1, l2, l3)



def run(initial, ν, λ, steps=2e3, step=1e-3):
    """Run Calculator with given parameters, and return a tuple with (t, theta, phi, lv, ls, lorth) arrays."""
    calc = Calculator(initial, ν, λ, step=step)
    zeros = lambda: np.zeros((steps))
    (t, theta, phi) = (zeros(), zeros(), zeros())
    (lv, ls, lorth) = (zeros(), zeros(), zeros())
    (l1, l2, l3) = (zeros(), zeros(), zeros())
    for (i, p) in enumerate(itertools.islice(calc, int(steps))):
        t[i] = i*step
        theta[i] = p.x
        phi[i] = p.phi
        (lv[i], ls[i], lorth[i]) = calc.spins(state=p)
        (l1[i], l2[i], l3[i]) = calc.orth_spins(state=p)
    return (t, theta, phi, Spins(lv, ls, lorth), StdSpins(l1, l2, l3))

def t_plot(t, theta, phi):
    fig = plt.figure(dpi=100)
    thaxes = fig.add_subplot(111)
    thaxes.set_ylabel("θ")
    thaxes.xaxis.set_label_text("t / s")
    phiaxes = fig.add_subplot(111, sharex=thaxes)
    phiaxes.yaxis.tick_right()
    phiaxes.yaxis.set_label_position("right")
    phiaxes.yaxis.set_label_text("ɸ")

    thaxes.plot(t, theta)
    phiaxes.plot(t, phi)
    plt.show()

def polar_axis_plot(theta, phi, axes=None, **kwargs):
    """Plot a bird's perspective of a spinner."""
    # Transform theta into polar argument r.
    r = np.sin(theta)

    if not axes:
        fig = plt.figure(dpi=100)
        axes = fig.add_subplot(111, projection='polar')
    axes.plot(phi, r, **kwargs)

def plot_assignment_parameters(steps=5000, step=1e-2):
    # 3 b(i)
    fig = plt.figure(figsize=(4 * 22/2.54, 76/2.54))

    theta = 0.4
    nu = 0.05
    initial = MoveParams(theta, 0, 0, 0)
    for i, lmbda in enumerate([math.cos(theta), 1.1*math.cos(theta), 1.4*math.cos(theta)]):
        (t, th, ph, (lv, ls, lorth), (l1, l2, l3)) = run(initial, nu, lmbda, steps=steps, step=step)

        axes = fig.add_subplot(4,4,1 +  4 * i, projection='polar')
        axes.set_ybound(0, 2)
        axes.set_title('Axis movement with λ = {:.2f}'.format(lmbda))
        polar_axis_plot(th, ph, axes=axes, label='λ = {:.2f}'.format(lmbda))

        # 3 (c)
        spinaxes = fig.add_subplot(4,4,2 + 4 * i)
        spinaxes.set_title('Angular momentum with λ = {:.2f}'.format(lmbda))
        spinaxes.plot(t, lv, label='l_v')
        spinaxes.plot(t, ls, label='l_s')
        spinaxes.plot(t, lorth, label='l_orth')
        spinaxes.legend()

        # 3 (d)
        orthspinaxes = fig.add_subplot(4,4,3 + 4 * i)
        orthspinaxes.set_title('Angular momentum in standard base with λ = {:.2f}'.format(lmbda))
        orthspinaxes.plot(t, l1, label='l1')
        orthspinaxes.plot(t, l2, label='l2')
        orthspinaxes.plot(t, l3, label='l3')
        orthspinaxes.legend()

        # 3 (e)
        xyaxes = fig.add_subplot(4,4,4 + 4 * i)
        xyaxes.set_title('Angular momentum direction with λ = {:.2f}'.format(lmbda))
        xyaxes.plot(l1, l2, label='l')
        xyaxes.legend()

    # 3 b(ii)
    axes = fig.add_subplot(4,4,13, projection='polar')
    axes.set_title('Axis movement (ii)')
    lmbda = 1.25
    theta = 0.664
    nu = 0.05
    initial = MoveParams(theta, 0, 0, 0)
    (t, th, ph, (lv, ls, lorth), (l1, l2, l3)) = run(initial, nu, lmbda, steps=steps, step=step)
    polar_axis_plot(th, ph, axes=axes, label='λ = {:.2f}'.format(lmbda))

    # 3 (c)
    spinaxes = fig.add_subplot(4,4,14)
    spinaxes.set_title('Spins with λ = {:.2f}'.format(lmbda))
    spinaxes.plot(t, lv, label='l_v')
    spinaxes.plot(t, ls, label='l_s')
    spinaxes.plot(t, lorth, label='l_orth')
    spinaxes.legend()

    # 3 (d)
    orthspinaxes = fig.add_subplot(4,4,15)
    orthspinaxes.set_title('Spins in standard base with λ = {:.2f}'.format(lmbda))
    orthspinaxes.plot(t, l1, label='l1')
    orthspinaxes.plot(t, l2, label='l2')
    orthspinaxes.plot(t, l3, label='l3')
    orthspinaxes.legend()

    # 3 (e)
    xyaxes = fig.add_subplot(4,4,16)
    xyaxes.set_title('Angular momentum direction with λ = {:.2f}'.format(lmbda))
    xyaxes.plot(l1, l2, label='l')
    xyaxes.legend()

    fig.legend(())
    fig.savefig('plot.pdf')