#!/usr/bin/env python3

import numpy as np
import matplotlib.pyplot as plt
from 	matplotlib.animation 	import FuncAnimation
from    matplotlib              import cm
import pandas as pd
import yaml
import os
import matplotlib.gridspec as gridspec
from matplotlib import rcParams
from matplotlib.ticker import MultipleLocator

class plotter_v3:
    # This code will make 1 figure with 2 subplots with a result as a function of time: e.g. T(t) or Fr_peak(t).
    # Subplot 1 is generally used for the ramp up.
    # Subplot 2 is generally used for the quench.

    # paths
    input_path_1 = []
    input_path_2 = []
    parameter = None

    # Figures and axes
    fig_square = True
    fig = None
    gs = None
    ax1 = None
    ax2 = None
    plot_every= 1

    max_values = []

    # Data
    input_data = None
    limits = None

    time1 = []
    time2 = []

    max_values1 = []
    max_values2 = []

    maxlim = -1e99
    minlim = 1e99

    def __init__(self, input_1, input_2=None, labels1=None, labels2=None, parameter=None):
        self.input_path_1 = input_1
        self.input_path_2 = input_2
        self.parameter = parameter

        if self.fig_square:
            # Figures and axes
            # self.fig = plt.figure(figsize=(12, 4))
            self.fig = plt.figure(figsize=(12, 3.5))
            self.gs = gridspec.GridSpec(1, 2, width_ratios=[1, 2])
            self.ax1 = self.fig.add_subplot(self.gs[0])
            self.ax2 = self.fig.add_subplot(self.gs[1])
        else:
            # Figures and axes
            self.fig = plt.figure(figsize=(6, 4))
            self.gs = gridspec.GridSpec(1, 2, width_ratios=[1, 1])
            self.ax1 = self.fig.add_subplot(self.gs[0])
            self.ax2 = self.fig.add_subplot(self.gs[1])

        # Making the figures:
        # https://matplotlib.org/stable/gallery/subplots_axes_and_figures/broken_axis.html
        # self.fig, (self.ax1, self.ax2) = plt.subplots(1, 2, sharey=True)
        self.fig.subplots_adjust(wspace=0.035)

        # plot the same data on both axes
        self.load_yaml()

        markers = ['o', 'd', '*', 's', 'v', '^']

        for i in range(len(self.time1)):
            if self.parameter in ['Fr', 'Fz', 'II', 'IR']:
                self.ax1.plot(self.time1[i], np.array(self.max_values1[i])*1e-3, marker=markers[i], mfc='white', label=labels1[i])
            elif self.parameter in ['Fr_density', 'Fz_density']:
                self.ax1.plot(self.time1[i], np.array(self.max_values1[i])*1e-9, marker=markers[i], mfc='white', label=labels1[i])
            else:
                self.ax1.plot(self.time1[i], np.array(self.max_values1[i]), marker=markers[i], mfc='white', label=labels1[i])
            if np.max(self.max_values1[i]) > self.maxlim: self.maxlim = np.max(self.max_values1[i])
            if np.min(self.max_values1[i]) > self.minlim: self.minlim = np.min(self.max_values1[i])

        for i in range(len(self.time2)):
            if self.parameter in ['Fr', 'Fz', 'II', 'IR']:
                self.ax2.plot(np.array(self.time2[i])-self.time2[i][0]+5000.0, np.array(self.max_values2[i])*1e-3, marker=markers[i], mfc='white', label=labels2[i])
            elif self.parameter in ['Fr_density', 'Fz_density']:
                self.ax2.plot(np.array(self.time2[i])-self.time2[i][0]+5000.0, np.array(self.max_values2[i])*1e-9, marker=markers[i], mfc='white', label=labels2[i])
            else:
                self.ax2.plot(np.array(self.time2[i])-self.time2[i][0]+5000.0, np.array(self.max_values2[i]), marker=markers[i], mfc='white', label=labels2[i])
            if np.max(self.max_values2[i]) > self.maxlim: self.maxlim = np.max(self.max_values2[i])
            if np.min(self.max_values2[i]) > self.minlim: self.minlim = np.min(self.max_values2[i])

        # zoom-in / limit the view to different portions of the data
        if self.parameter in ['Fr', 'Fz', 'II', 'IR']:
            self.ax1.set_ylim(0, self.maxlim * 1.1 * 1e-3)
            self.ax2.set_ylim(0, self.maxlim * 1.1 * 1e-3)
        elif self.parameter in ['Fr_density', 'Fz_density']:
            self.ax1.set_ylim(0, self.maxlim * 1.1 * 1e-9)
            self.ax2.set_ylim(0, self.maxlim * 1.1 * 1e-9)
        else:
            self.ax1.set_ylim(0, self.maxlim*1.1)
            self.ax2.set_ylim(0, self.maxlim*1.1)
        # self.ax1.set_xlim(0, 5107.35)  # most of the data

        self.ax1.set_xlim(0, 5000)  # most of the data
        self.ax2.set_xlim(5000, 5002)  # most of the data

        self.ax1.set_xlabel('Time, s')
        self.ax2.set_xlabel('Time, s')

        if self.parameter in ['Fr', 'Fz']:
            self.ax1.set_ylabel('Peak Force, kN')
        elif self.parameter in ['Fr_density', 'Fz_density']:
            self.ax1.set_ylabel('Peak Force Density, N/mm$^3$')
        elif self.parameter in ['I_density']:
            self.ax1.set_ylabel('Peak Current Density, A/mm$^2$')
        elif self.parameter in ['II', 'IR']:
            self.ax1.set_ylabel('Current, kA')
        elif self.parameter == 'T':
            self.ax1.set_ylabel('Temperature, K')
        elif self.parameter in ['Bz', 'Br']:
            self.ax1.set_ylabel('Magnetic Field, T')

        if self.parameter in ['T']:
            self.ax1.legend(loc='upper left')
            self.ax2.legend(loc='upper left')
        else:
            self.ax1.legend(loc='lower right')
            self.ax2.legend(loc='upper right')

        self.ax1.grid()
        self.ax2.grid()


        # hide the spines between ax and ax2
        self.ax1.spines.right.set_visible(False)
        self.ax2.spines.left.set_visible(False)
        self.ax1.xaxis.tick_bottom()
        self.ax1.tick_params(labeltop=False)  # don't put tick labels at the top
        self.ax2.xaxis.tick_bottom()
        self.ax2.tick_params(left = False, labelleft = False)

        # Now, let's turn towards the cut-out slanted lines.
        # We create line objects in axes coordinates, in which (0,0), (0,1),
        # (1,0), and (1,1) are the four corners of the axes.
        # The slanted lines themselves are markers at those locations, such that the
        # lines keep their angle and position, independent of the axes size or scale
        # Finally, we need to disable clipping.

        d = 2.0  # proportion of vertical to horizontal extent of the slanted line
        kwargs = dict(marker=[(-1, -d), (1, d)], markersize=12,
                      linestyle="none", color='k', mec='k', mew=1, clip_on=False)
        self.ax1.plot([1, 1], [0, 1], transform=self.ax1.transAxes, **kwargs)
        self.ax2.plot([0, 0], [0, 1], transform=self.ax2.transAxes, **kwargs)


        plt.setp(self.ax1.get_xticklabels()[-1], visible=False)

        for axis in ['top', 'bottom', 'left']:
            self.ax1.spines[axis].set_linewidth(1.0)

        for axis in ['top', 'bottom', 'right']:
            self.ax2.spines[axis].set_linewidth(1.0)

        self.ax1.xaxis.set_tick_params(direction='in', top='on', width=1.2, length=5)
        self.ax1.yaxis.set_tick_params(direction='in', left='on', width=1.2, length=5)
        self.ax1.xaxis.set_tick_params(direction='in', top='on', which='minor', width=1.0, length=3)
        self.ax1.yaxis.set_tick_params(direction='in', left='on', which='minor', width=1.0, length=3)

        self.ax2.xaxis.set_tick_params(direction='in', top='on', width=1.2, length=5)
        self.ax2.yaxis.set_tick_params(direction='in', right='on', width=1.2, length=5)
        self.ax2.xaxis.set_tick_params(direction='in', top='on', which='minor', width=1.0, length=3)
        self.ax2.yaxis.set_tick_params(direction='in', right='on', which='minor', width=1.0, length=3)



        plt.show()

        # self.load_yaml()

        # self.make_gif()

    def load_yaml(self):
        # Reads the contents of the yaml file
        for i in self.input_path_1:
            with open(i + os.sep + 'Paraview' + os.sep + 'python_files_to_load_' + self.parameter + '.yaml') as file:
                data = yaml.load(file, Loader=yaml.FullLoader)
                self.time1.append(data['time'])
                self.max_values1.append(data['max_value'])

        for i in self.input_path_2:
            with open(i + os.sep + 'Paraview' + os.sep + 'python_files_to_load_' + self.parameter + '.yaml') as file:
                data = yaml.load(file, Loader=yaml.FullLoader)
                self.time2.append(data['time'])
                self.max_values2.append(data['max_value'])

    def make_gif(self):

        self.time_steps = list(self.input_data.keys())
        step_data = self.input_data[self.time_steps[4]]

        # self.limits['minmax'][1] = 0.5 * self.limits['minmax'][1]

        for i in step_data:
            df = pd.read_csv(i)
            self.ax.tricontourf(df['x'].values, df['y'].values, df['v'].values, vmin=self.limits['minmax'][0], vmax=self.limits['minmax'][1], cmap=cm.coolwarm)
            self.ax.tricontourf(-1*df['x'].values, df['y'].values, df['v'].values, vmin=self.limits['minmax'][0], vmax=self.limits['minmax'][1], cmap=cm.coolwarm)

        self.ax.set_aspect('equal', 'box')

        anim = FuncAnimation(self.fig, self.update_gif, frames=np.arange(1, len(self.time_steps), self.plot_every), interval=100)
        # IF YOU GET AN ERROR IN SAVING THE GIF -> DOWNLOAD IMAGEMAGICK
        # anim.save(fname, dpi=120, writer='imagemagick', fps=20)
        fname = self.rpath + os.sep + 'test.gif'
        anim.save(fname, dpi=300, fps=10)
        plt.close('all')

    def update_gif(self, i):
        # Updates the gif
        self.ax.clear()

        print('Frame', i)

        step_data = self.input_data[self.time_steps[i]]

        for j in step_data:
            df = pd.read_csv(j)
            self.ax.tricontourf(df['x'].values, df['y'].values, df['v'].values, vmin=self.limits['minmax'][0], vmax=self.limits['minmax'][1], cmap=cm.coolwarm)
            self.ax.tricontourf(-1*df['x'].values, df['y'].values, df['v'].values, vmin=self.limits['minmax'][0], vmax=self.limits['minmax'][1], cmap=cm.coolwarm)

        self.ax.set_title('t = ' + str(np.round((self.time_steps[i]-self.time_steps[0]) * 1e3)) + ' ms')
        self.ax.set_aspect('equal', 'box')

if __name__ == "__main__":
    # input_1 = ['C:/Tim Mulder/CERNBox/Staff - TE-MPE-PE/Projects/ni-coils/geometry/output/test_geometry_4_MT_output/results/result - 2023-08-29 - 0.1 - 2']
    # # input_2 = ['C:/Tim Mulder/CERNBox/Staff - TE-MPE-PE/Projects/ni-coils/geometry/output/test_geometry_4_MT_output/results/result - 2023-09-08 - 0.1 - 5',
    # #            'C:/Tim Mulder/CERNBox/Staff - TE-MPE-PE/Projects/ni-coils/geometry/output/test_geometry_4_MT_output/results/result - 2023-08-29 - 0.1 - 3']
    # # labels2 = ['CD Full', 'Natural Extr.']
    # input_2 = ['C:/Tim Mulder/CERNBox/Staff - TE-MPE-PE/Projects/ni-coils/geometry/output/test_geometry_4_MT_output/results/result - 2023-09-08 - 0.1 - 5',
    #            'C:/Tim Mulder/CERNBox/Staff - TE-MPE-PE/Projects/ni-coils/geometry/output/test_geometry_4_MT_output/results/result - 2023-09-04 - 0.1 - 2',
    #            'C:/Tim Mulder/CERNBox/Staff - TE-MPE-PE/Projects/ni-coils/geometry/output/test_geometry_4_MT_output/results/result - 2023-08-29 - 0.1 - 6',
    #            'C:/Tim Mulder/CERNBox/Staff - TE-MPE-PE/Projects/ni-coils/geometry/output/test_geometry_4_MT_Cu1_output/results/result - 2023-09-05 - 0.1 - 2',
    #            'C:/Tim Mulder/CERNBox/Staff - TE-MPE-PE/Projects/ni-coils/geometry/output/test_geometry_4_MT_output/results/result - 2023-08-29 - 0.1 - 3',]
    # labels1 = ['Ramp up']
    # labels2 = ['1) CD Full', '2) CD Center', '3) CD Extremities', '4) CD Center + QB', '5) Natural Extr. 0.5']


    input_1 = ['C:/Tim Mulder/CERNBox/Staff - TE-MPE-PE/Projects/ni-coils/geometry/output/test_geometry_23_2_output/results/result - 2023-10-05 - 0.1 - 1',
               'C:/Tim Mulder/CERNBox/Staff - TE-MPE-PE/Projects/ni-coils/geometry/output/test_geometry_23_3_output/results/result - 2023-10-06 - 0.1 - 1',
               'C:/Tim Mulder/CERNBox/Staff - TE-MPE-PE/Projects/ni-coils/geometry/output/test_geometry_23_5_output/results/result - 2023-10-13 - 0.1 - 1',
               'C:/Tim Mulder/CERNBox/Staff - TE-MPE-PE/Projects/ni-coils/geometry/output/test_geometry_23_6_output/results/result - 2023-10-23 - 0.1 - 1',
               'C:/Tim Mulder/CERNBox/Staff - TE-MPE-PE/Projects/ni-coils/geometry/output/test_geometry_23_4_output/results/result - 2023-10-10 - 0.1 - 1',]
    input_2 = ['C:/Tim Mulder/CERNBox/Staff - TE-MPE-PE/Projects/ni-coils/geometry/output/test_geometry_23_2_output/results/result - 2023-10-09 - 0.1 - 1',
               'C:/Tim Mulder/CERNBox/Staff - TE-MPE-PE/Projects/ni-coils/geometry/output/test_geometry_23_3_output/results/result - 2023-10-09 - 0.1 - 1']
    labels1 = ['W/ screening currents', 'W/o screening currents', 'W/ SC 4 mm extr. coils', 'W/ SC 4 mm extr. coils v2',  'W/o SC 4 mm extr. coils']
    labels2 = ['With screening currents', 'Without screening currents']

    # input_1 = ['C:/Tim Mulder/CERNBox/Staff - TE-MPE-PE/Projects/ni-coils/geometry/output/test_geometry_28_FCS/results/result - 2024-05-22 - 0.1 - 1']
    # input_2 = ['C:/Tim Mulder/CERNBox/Staff - TE-MPE-PE/Projects/ni-coils/geometry/output/test_geometry_4_MT_output/results/result - 2023-09-08 - 0.1 - 5',
    #            'C:/Tim Mulder/CERNBox/Staff - TE-MPE-PE/Projects/ni-coils/geometry/output/test_geometry_4_MT_output/results/result - 2023-08-29 - 0.1 - 3']
    # labels2 = ['CD Full', 'Natural Extr.']
    # input_2 = ['C:/Tim Mulder/CERNBox/Staff - TE-MPE-PE/Projects/ni-coils/geometry/output/test_geometry_28_FCS/results/result - 2024-05-28 - 0.1 - 1',
    #            'C:/Tim Mulder/CERNBox/Staff - TE-MPE-PE/Projects/ni-coils/geometry/output/test_geometry_28_FCS/results/result - 2024-05-28 - 0.1 - 2', ]
    # labels1 = ['Ramp up']
    # labels2 = ['1) Slow Protection', '2) Faster Protection']


    ptr = plotter_v3(input_1, input_2, labels1, labels2, 'Fr_density')