import scinet
import data_gen

import argparse

import tqdm
import torch
import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
from matplotlib.ticker import FuncFormatter, MultipleLocator


def _set_pi_ticks(*axis_list):
    '''Format pyplot axis ticks to use multiples of π'''

    for axis in axis_list:

        axis.set_major_formatter(FuncFormatter(
            lambda val, pos:
            '{:.0g}$\pi$'.format(val / np.pi) if val else '0'
        ))
        axis.set_major_locator(MultipleLocator(base=np.pi))


def visualize_sample(obs_θ, obs_φ, model_θ, model_φ):
    '''Visualize and compare the Earth angles (θ) from observations and model

    Creates an animated figure of both polar and cartesian plots of the
    Earth angles θ_e and θ_m, over the entire time evolution.

    From the input datasets a random orbit (N) will be chosen for visualization.

    Parameters
    ----------
    obs_θ : numpy.ndarray
        Array of observation (simulated) Earth angles

    obs_φ : numpy.ndarray
        Array of observation (simulated) Sun angles

    model_θ : numpy.ndarray
        Array of Earth angles generated by the trained network

    model_φ : numpy.ndarray
        Array of Sun angles generated by the trained network

    '''

    import matplotlib.animation as anim

    def _update_plot(n):
        '''Update plot lines for animation'''

        lines[0].set_data(obs_θ[:n, 0], R[:n])
        lines[1].set_data(model_θ[:n, 0], R[:n])

        lines[2].set_data(obs_θ[:n, 1], R[:n])
        lines[3].set_data(model_θ[:n, 1], R[:n])

        lines[4].set_data(x[:n], obs_θ[:n, 0])
        lines[5].set_data(x[:n], model_θ[:n, 0])
        # lines[6].set_data(x[:n], diff_θ[:n, 0])

        lines[6].set_data(x[:n], obs_θ[:n, 1])
        lines[7].set_data(x[:n], model_θ[:n, 1])
        # lines[9].set_data(x[:n], diff_θ[:n, 1])

        return lines

    # ----------------------------------------------------------------------
    # Select a random orbit to visualize.
    # The first observation datapoint (in the time series) is also ignored,
    # as there is not corresponding model output
    # ----------------------------------------------------------------------

    rand_ind = np.random.randint(obs_θ.shape[0])

    obs_θ = obs_θ[rand_ind, 1:]
    obs_φ = obs_φ[rand_ind, 1:]
    model_θ = model_θ[rand_ind, :]
    model_φ = model_φ[rand_ind, :]

    # diff_θ = obs_θ - model_θ

    R = np.ones_like(obs_θ[:, 0])
    x = range(R.size)

    # ----------------------------------------------------------------------
    # Setup figure
    # ----------------------------------------------------------------------

    fig = plt.figure()
    fig.suptitle(f'Sample orbit {rand_ind}')

    ax1, ax2 = fig.add_subplot(221, polar=1), fig.add_subplot(222, polar=1)
    ax3, ax4 = fig.add_subplot(223), fig.add_subplot(224)

    ax1.set_title(r'Sun Angle ($θ_e$)')
    ax2.set_title(r'Mars Angle ($θ_m$)')

    _set_pi_ticks(ax3.yaxis, ax4.yaxis)

    ax1.plot(0, 0, 'go')
    ax2.plot(0, 0, 'go')

    # ----------------------------------------------------------------------
    # Plot datasets
    # ----------------------------------------------------------------------

    [line1_o] = ax1.plot(obs_θ[:, 0], R, '-go', markevery=[-1])
    [line1_m] = ax1.plot(model_θ[:, 0], R, '-ro', markevery=[-1])

    [line2_o] = ax2.plot(obs_θ[:, 1], R, '-go', markevery=[-1])
    [line2_m] = ax2.plot(model_θ[:, 1], R, '-ro', markevery=[-1])

    [line3_o] = ax3.plot(x, obs_θ[:, 0], '-go')
    [line3_m] = ax3.plot(x, model_θ[:, 0], '-ro')
    # [line3_d] = ax3.plot(x, diff_θ[:, 0], '-bo')

    [line4_o] = ax4.plot(x, obs_θ[:, 1], '-go')
    [line4_m] = ax4.plot(x, model_θ[:, 1], '-ro')
    # [line4_d] = ax4.plot(x, diff_θ[:, 1], '-bo')

    lines = [line1_o, line1_m, line2_o, line2_m,
             line3_o, line3_m, line4_o, line4_m]

    fig.legend(lines, ['Observations', 'Model'])

    # ----------------------------------------------------------------------
    # Animate the figure
    # ----------------------------------------------------------------------

    anim.FuncAnimation(fig, _update_plot, model_θ.shape[0], blit=True)

    plt.show()


class TimeEvolvedScinet(scinet.Scinet):
    '''Subclass of `Scinet` with additional time-evolution of latent layers

    For the `copernicus` example, we wish to examine the time-evolution
    capabilities of `Scinet`, which requires a slight alteration of the
    network architecture.

    Time evolution is inserted into the VAE in the form of a small recurrent
    neural network (τ), which acts on the latent neurons, in between the
    encoding and decoding layers.

    Specifically for the `copernicus` example, the RNN takes the form of an
    exceedingly simply map: r(t) -> r(t) + b

    The τ bias is constant throughout the network, and is designed to
    force the learning of the simplest possible time evolution.

    See the `train` and `forward` methods for more details

    Parameters
    ----------
    hyp : hyperparameters.Hyperparameters
        Hyperparameters object guiding the setup of the network

    Attributes
    ----------
    τ_bias : torch.nn.Parameter
        A tensor attached to the network optimizer, containing the additive
        time evolution bias (b)

    '''

    def __init__(self, hyp):
        super().__init__(hyp)

        N = hyp.latentNodes

        self.τ_bias = torch.nn.Parameter(torch.zeros(N, requires_grad=True))
        self.τ_layers = 3

        self.optimizer = hyp.optimizer(self.parameters(), hyp.learningRate)

    def forward(self, x, first_pass=False):
        '''Network forward function

        Variational auto-encoder based forward function, with encoding,
        reparamaterization and decoding.

        Notes
        -----
        Differs from the base `Scinet` forward function by including a check
        for whether encoding is necessary, and the pass through the time
        evolution network (τ). Also does not require a question neuron.

        Parameters
        ----------
        x : torch.Tensor
            Input training dataset batch

        first_pass : bool, optional
            Whether this iteration of training requires encoding or only time
            evolution and decoding. Corresponds to the training of the RNN
            layers

        Returns
        -------
        answer : torch.Tensor
            The time evolved, then decoded final VAE neurons

        evolved_latent : torch.Tensor
            The time evolved but not decoded latent neurons

        μ : torch.Tensor
            Latent neuron reparameterization mean

        σ : torch.Tensor
            Latent neuron reparameterization standard deviation

        '''

        μ, σ = None, None

        # Pass through encoder layers, if the input is not a latent layer
        if first_pass:

            # Pass through encoder
            μ, σ = self.encode(x)

            # Reparamaterize
            x = self.reparameterize(μ, σ)

        # Pass through the time evolution (τ) network
        for _ in range(self.τ_layers):
            x = x + self.τ_bias

        evolved_latent = x

        # Pass through decoder
        answer = self.decode(x, torch.tensor([]))

        return answer, evolved_latent, μ, σ

    def train(self, observations, batch_N):
        '''Network training function

        Batched training function, splitting up the input training dataset into
        batches before passing to the forward function and performing
        back-propogation.

        Notes
        -----
        Differs from the base `Scinet` training function by including an inner
        loop over the time series, and a corresponding adjustment of the
        forward and backward functions based on the position in the network.

        Parameters
        ----------
        observations : torch.Tensor
            Input training dataset

        batch_N : int
            Batch size

        Returns
        -------
        avgLoss : float
            The averaged value of the cost function over all batches and τ
        '''

        avgLoss = 0
        N = observations.shape[0]

        # Split datset into batches
        for i in range(0, N, batch_N):

            obs_batch = observations[i:i + batch_N].to(self.device)

            # initial model input is given by first observation
            obs = obs_batch[:, 0, :]

            # iterate over time series
            for te_step in range(observations.shape[1] - 1):
                first_pass = not te_step

                # Target θ given by proceeding observation in time series
                ans = obs_batch[:, te_step + 1, :]

                self.zero_grad()

                # Future model input given by evolved latent layers
                outputs, obs, μ, σ = self(obs, first_pass)

                # Loss function and back-propogation
                if first_pass:
                    loss = self.lossFunct(μ, σ, outputs, ans)
                else:
                    loss = self.leadingLoss(outputs, ans)

                loss.backward(retain_graph=True)

                self.optimizer.step()

                # Average losses
                avgLoss += loss.item() / observations.shape[1]

            self.trainCounter += 1

        avgLoss /= (N // batch_N)

        return avgLoss

    def test(self, observations):
        '''Network testing function

        Test the model effectiveness by passing the input observations through
        the trained network, to observe the final accuracy and the activations
        of the final latent layers.

        Identical to `train`, with the exclusion of all back-propogation,
        gradient descent and batching.

        Notes
        -----
        Differs from the base `Scinet` testing function by including an inner
        loop over the time series, and a corresponding adjustment of the
        forward and cost functions based on the position in the network.

        Parameters
        ----------
        observations : torch.Tensor
            Input testing dataset

        Returns
        -------
        avgLoss : float
            The averaged value of the cost function over all τ

        model_θ : torch.Tensor
            Output of trained network answer neurons

        model_φ : torch.Tensor
            Activation of trained network latent neurons
        '''

        avgLoss = 0

        observations = observations.to(self.device)
        model_θ = np.empty_like(observations)
        model_φ = np.empty_like(observations)

        with torch.no_grad():

            # initial forward function input is given by observation
            obs = observations[:, 0, :]

            # iterate over time series
            for te_step in range(observations.shape[1] - 1):
                first_pass = not te_step

                # Target θ given by proceeding observation in time series
                ans = observations[:, te_step + 1, :]

                # Future model input given by evolved latent layers
                model_ans, obs, μ, σ = self(obs, first_pass)

                model_θ[:, te_step + 1, :] = model_ans
                model_φ[:, te_step + 1, :] = obs

                # Loss functions
                if first_pass:
                    loss = self.lossFunct(μ, σ, model_ans, ans)
                else:
                    loss = self.leadingLoss(model_ans, ans)

                avgLoss += loss.item()

            avgLoss /= observations.shape[1]

        # The first of the time series is not decoded before passing through τ
        return avgLoss, model_θ[:, 1:, :], model_φ[:, 1:, :]


if __name__ == '__main__':

    # ----------------------------------------------------------------------
    # Parse command line arguments
    # ----------------------------------------------------------------------

    parser = argparse.ArgumentParser()

    parser.add_argument('-M', type=int, default=5,
                        help="Number of time-evolution steps")
    parser.add_argument('-N', type=int, default=15000,
                        help="Training dataset size")
    parser.add_argument('--test-N', type=int, default=1000,
                        help="Testing dataset size")
    parser.add_argument('-t', '--del-t', type=int, default=7,
                        help="Number of days per time-evolution step")

    parser.add_argument('-a', type=float, default=0.001,
                        help="Learning rate (α)")
    parser.add_argument('-b', type=int, default=2000,
                        help="Training batch size")
    parser.add_argument('-E', type=int, default=25,
                        help="Number of training epochs")
    parser.add_argument('-B', '--beta-anneal', action='store_true',
                        help="Whether to use 'Beta-Annealing'")

    parser.add_argument('-v', '--verbose', action='store_true')
    parser.add_argument('--plot-loss', action='store_true',
                        help='Show a plot of loss over training epochs')

    args = parser.parse_args()

    hyp = scinet.Hyperparameters()
    hyp.learningRate = args.a
    hyp.encoderNodes[0] = 2
    hyp.latentNodes = 2
    hyp.decoderNodes[-1] = 2
    hyp.questionNodes = 0
    hyp.leadingLoss = torch.nn.MSELoss
    if args.beta_anneal:
        hyp.annealEpoch = 1e8

    N = args.N
    M = args.M
    del_t = args.del_t
    test_N = args.test_N

    batch_size = args.b
    epochs = tqdm.tqdm(range(args.E), disable=(not args.verbose),
                       bar_format="{l_bar}{bar}| Epoch {n_fmt}/{total_fmt} ["
                                  "{elapsed}<{remaining}, {rate_fmt}{postfix}]")

    # ----------------------------------------------------------------------
    # Initialize network and generate training data
    # ----------------------------------------------------------------------

    model = TimeEvolvedScinet(hyp)

    φ_e, φ_m, θ_e, θ_m = data_gen.generate_orbits(N, M, del_t)

    # data_gen.anim_orbit(φ_e, φ_m, θ_e, θ_m)

    # ----------------------------------------------------------------------
    # Train the model
    # ----------------------------------------------------------------------

    loss = []
    for _ in epochs:

        # ------------------------------------------------------------------
        # Observations are all the elements of θ
        # ------------------------------------------------------------------

        obs = torch.from_numpy(np.concatenate(
            (θ_e[..., np.newaxis], θ_m[..., np.newaxis]), axis=-1
        )).float()

        # ------------------------------------------------------------------
        # Train the network, record the average losses for each epoch
        # ------------------------------------------------------------------

        loss.append(model.train(obs, batch_size))

        if args.verbose:
            epochs.set_description_str(f"loss: {loss[-1]:.2f}")

    if args.plot_loss:
        scinet.plot_loss(loss)

    # ----------------------------------------------------------------------
    # Generate testing data
    # ----------------------------------------------------------------------

    φ_et, φ_mt, θ_et, θ_mt = data_gen.generate_orbits(test_N, M, del_t, True)

    test_θ = torch.from_numpy(np.concatenate(
        (θ_et[..., np.newaxis], θ_mt[..., np.newaxis]), axis=-1
    )).float()

    test_φ = torch.from_numpy(np.concatenate(
        (φ_et[..., np.newaxis], φ_mt[..., np.newaxis]), axis=-1
    )).float()

    # ----------------------------------------------------------------------
    # Test the trained network
    # ----------------------------------------------------------------------

    loss, out_θ, out_φ = model.test(test_θ)

    # ----------------------------------------------------------------------
    # Plot the results (latent activation) of the test
    # ----------------------------------------------------------------------

    # visualize_sample(test_θ, test_φ, out_θ, out_φ)

    plot_kwargs = {"method": 'surface', "axlabels": ['φ_e', 'φ_m'],
                   "axis_formatter": _set_pi_ticks}

    space = np.linspace(0, 2 * np.pi, num=50)
    mesh_e, mesh_m = np.meshgrid(space, space)

    scinet.plot_latent(mesh_e, mesh_m, out_φ[:, -1, :], **plot_kwargs)

    plt.show()
