#!/usr/bin/env python3

import numpy as np
import matplotlib.pyplot as plt
from 	matplotlib.animation 	import FuncAnimation
# from    matplotlib              import cm
import matplotlib
import pandas as pd
from scipy import integrate

import os
import matplotlib.gridspec as gridspec
from matplotlib import rcParams
from matplotlib.ticker import MultipleLocator

rcParams['axes.labelsize'] =11
rcParams['xtick.labelsize'] =11
rcParams['ytick.labelsize'] =11
rcParams['legend.fontsize'] =11
rcParams['font.sans-serif'] = 'Arial'
rcParams['text.usetex'] = False
rcParams['figure.autolayout'] = True
rcParams['lines.markersize'] = 5
rcParams['figure.figsize'] = (7, 6)
rcParams['axes.grid'] = True
rcParams['grid.color'] = 'lightGrey'
rcParams['axes.axisbelow'] = True

class plotter_v3:
    # This plotter will make 2d plots of up to two results folders

    # paths
    input_paths = []
    parameter = None

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

    max_values = []

    # Data
    input_data = None
    limits = None

    time1 = []
    time2 = []

    max_values1 = []
    max_values2 = []

    maxlim = -1e99
    minlim = 1e99
    xlim = None

    def __init__(self, input_folder, xlim = None, parameter=None):
        self.input_paths = input_folder
        self.parameter = parameter
        self.xlim = xlim

        if self.fig_square:
            # Figures and axes
            self.fig = plt.figure(figsize=(12, 3.5))
            self.ax1 = self.fig.add_subplot()
        else:
            # Figures and axes
            # self.fig = plt.figure(figsize=(8, 4))
            self.fig = plt.figure(figsize=(6,5))
            self.fig2 = plt.figure(figsize=(10,5))
            self.fig3 = plt.figure(figsize=(6,5))
            self.ax1 = self.fig.add_subplot(111)
            self.ax2 = self.fig2.add_subplot(111)
            self.ax3 = self.fig3.add_subplot(111)

        # 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)

        # Loading the data

        data = pd.read_csv(self.input_paths + os.sep + 'Power_per_Coil.csv')
        data2 = pd.read_csv(self.input_paths + os.sep + 'Power_per_Solenoid.csv')
        columns = data.columns
        columns2 = data2.columns

        # cmap = cm.get_cmap('viridis')
        cmap = matplotlib.colormaps['coolwarm']
        colors = [cmap(i) for i in np.linspace(0, 1.0, len(columns))]
        x = data['time'].values
        for i in range(len(columns)):
            if columns[i] != 'time':
                y = data[columns[i]].values
                self.ax1.plot(x, y*1e-6, color=colors[i])

        cmap = matplotlib.colormaps['coolwarm']
        colors2 = [cmap(i) for i in np.linspace(0, 1.0, len(columns2))]
        for i in range(len(columns2)):
            if columns2[i] != 'time':
                y = data2[columns2[i]].values
                self.ax2.plot(x, y*1e-6, color=colors2[i], label=columns2[i])
                self.ax3.plot(x, integrate.cumulative_trapezoid(y, x, initial=0), color=colors2[i], label=columns2[i])

        #

        self.ax1.set_xlabel('Time, s')
        self.ax2.set_xlabel('Time, s')
        self.ax1.set_ylabel('Power, MW')
        self.ax2.set_ylabel('Power, MW')
        if self.xlim is not None:
            self.ax1.set_xlim((self.xlim[0], self.xlim[1]))
            self.ax2.set_xlim((self.xlim[0], self.xlim[1]))

        self.ax2.legend()
        self.ax3.legend()

        # self.ax2.set_xlabel('Time, s')
        # self.ax2.set_xlim((-0.1, 0.8))

        # self.ax2.legend(loc='lower right', ncols=2)



        # self.ax1.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)

        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.ax1.xaxis.set_tick_params(direction='in', top='on', width=1.2, length=5)
        self.ax1.yaxis.set_tick_params(direction='in', right='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', right='on', which='minor', width=1.0, length=3)

        # self.ax1.xaxis.set_minor_locator(MultipleLocator(0.1))
        # self.ax1.xaxis.set_major_locator(MultipleLocator(0.2))

        # self.ax1.yaxis.set_minor_locator(MultipleLocator(10.0))
        # self.ax1.yaxis.set_major_locator(MultipleLocator(20.0))

        # 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)

        # self.ax2.xaxis.set_minor_locator(MultipleLocator(0.1))
        # self.ax2.xaxis.set_major_locator(MultipleLocator(0.2))

        # self.ax2.yaxis.set_minor_locator(MultipleLocator(25.0))
        # self.ax2.yaxis.set_major_locator(MultipleLocator(50.0))

        self.fig.tight_layout()
        self.fig2.tight_layout()
        # Poster
        self.fig.savefig(self.input_paths + os.sep + 'Power_per_element.png', dpi=1000)
        self.fig2.savefig(self.input_paths + os.sep + 'Power_per_coil.png', dpi=1000)



        plt.show()

if __name__ == "__main__":
    # inputs = ['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',
    #            '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 - 3',
    #           'C:/Tim Mulder/CERNBox/Staff - TE-MPE-PE/Projects/ni-coils/geometry/output/test_geometry_28_FCS/results/result - 2024-05-28 - 0.1 - 4',]
    # labels = ['1) 10 ms prot. delay', '2) 100 ms prot. delay', '3) No capacitor', '4) No protection']
    # inputs = ['C:/Tim Mulder/CERNBox/Staff - TE-MPE-PE/Projects/ni-coils/geometry/output/FCS_750/results/result - 2024-07-31 - 0.1 - 22',   # GOOD fast protection
    #             'C:/Tim Mulder/CERNBox/Staff - TE-MPE-PE/Projects/ni-coils/geometry/output/FCS_750/results/result - 2024-07-31 - 0.1 - 19',   # Good 300 V capacitor
    #           'C:/Tim Mulder/CERNBox/Staff - TE-MPE-PE/Projects/ni-coils/geometry/output/FCS_750/results/result - 2024-07-31 - 0.1 - 18',
    #           'C:/Tim Mulder/CERNBox/Staff - TE-MPE-PE/Projects/ni-coils/geometry/output/FCS_750/results/result - 2024-07-31 - 0.1 - 17',
    #           'C:/Tim Mulder/CERNBox/Staff - TE-MPE-PE/Projects/ni-coils/geometry/output/FCS_750/results/result - 2024-07-31 - 0.1 - 20',   # GOOD Only breaker
    #           'C:/Tim Mulder/CERNBox/Staff - TE-MPE-PE/Projects/ni-coils/geometry/output/FCS_750/results/result - 2024-07-31 - 0.1 - 21',   # GOOD No protection

              # 'C:/Tim Mulder/CERNBox/Staff - TE-MPE-PE/Projects/ni-coils/geometry/output/FCS_375/results/result - 2024-08-01 - 0.1 - 2',
              # 'C:/Tim Mulder/CERNBox/Staff - TE-MPE-PE/Projects/ni-coils/geometry/output/FCS_375/results/result - 2024-08-01 - 0.1 - 3',
              # 'C:/Tim Mulder/CERNBox/Staff - TE-MPE-PE/Projects/ni-coils/geometry/output/FCS_375/results/result - 2024-08-01 - 0.1 - 4',
              # 'C:/Tim Mulder/CERNBox/Staff - TE-MPE-PE/Projects/ni-coils/geometry/output/FCS_375/results/result - 2024-08-01 - 0.1 - 5',
              # 'C:/Tim Mulder/CERNBox/Staff - TE-MPE-PE/Projects/ni-coils/geometry/output/FCS_375/results/result - 2024-08-01 - 0.1 - 6',
              # 'C:/Tim Mulder/CERNBox/Staff - TE-MPE-PE/Projects/ni-coils/geometry/output/FCS_375/results/result - 2024-08-01 - 0.1 - 7',
              # 'C:/Tim Mulder/CERNBox/Staff - TE-MPE-PE/Projects/ni-coils/geometry/output/FCS_750/results/result - 2024-08-01 - 0.1 - 1',
              # ]
    # labels = ['CD 300V 50ms', 'CD 900V', 'CD 600V', 'CD 300V', 'Only breaker','No protection','7','8','9','10','11','12','13','14','15','16','17','18']
    # input_folder = 'C:/Tim Mulder/CERNBox/Staff - TE-MPE-PE/Projects/ni-coils/geometry/output/test_geometry_29_AMS_100_v2/results/result - 2024-07-03 - 0.1 - 5'
    input_folder = 'C:/Tim Mulder/CERNBox/Staff - TE-MPE-PE/Projects/ni-coils/geometry/output/test_geometry_29_AMS_100_v3/results/result - 2024-09-23 - 0.1 - 2'
    # input_folder = 'C:/Tim Mulder/CERNBox/Staff - TE-MPE-PE/Projects/ni-coils/geometry/output/FCS_750/results/result - 2024-07-31 - 0.1 - 21'

    ptr = plotter_v3(input_folder, xlim=[31.5, 34])