#!/usr/bin/env python3

import numpy as np
import pickle


# from scipy.io import loadmat
import os
import time
import shutil
import logging
import yaml
import sys
import itertools
# import 	telegram
from datetime import datetime
import multiprocessing
from multiprocessing import Pool
import  matplotlib.pyplot as plt
from    pathlib import Path
from    nicqs.lib import RegularGridInterpolator

# Torch
import torch

from 	steammaterials.STEAM_materials import STEAM_materials

# Scipy
import scipy
import scipy.integrate
from scipy import integrate
from scipy.interpolate import interp1d
from scipy.optimize import fsolve

class NI_solver_light:
    circuit_interp_dIdt = dict()
    t_array = []
    I_array = []
    I_density_array = []
    solver_t_array = []
    stored_energy = []
    V_array = []
    P_array = []
    Bz_center_array = []
    E_dis = None

    Power = []
    I_circuit = []
    V_circuit = []

    R_multiplier = 1.0

    unique_circuits = None
    include_Ic = True

    ## PID TEST
    max_rr = 10.0
    min_rr = 0.0
    setpoint = 200.0
    nominal_current = 760
    integral = 0.0
    Kp = 0.25
    Kd = 1.0
    Ki = 0.00001
    rr = 0.0
    pid_bool = False

    ## Saving
    savemod = 20
    save_n = 0

    ## Solver
    fast = True
    insulated = True

    # paths
    results_path = ''
    specific_output_path = ''
    sim_name = ''
    sim_output = ''
    file_summary = ''
    file_current = ''
    file_current_R = ''
    file_I_density = ''

    delimiter = ','
    export_geometry = False
    steam_mat_Ic = True
    SC_properties = {'fit-name': 'CFUN_HTS_JcFit_Fujikura_v1', 'n-value': 30, 'Ic': 200.0, 'Tref': 4.5, 'Bref': 40, 'angle': 0.0}

    def __init__(self, input_yaml, path_dict = None):
        # Loads the input yaml file.
        self.input_yaml_path = input_yaml
        self.path_dict = path_dict
        self.load_yaml()

        # Create folders and files
        self.prepare_files()

        # Prepares the logging and results file.
        self.create_ramp_rate_interpolation()

        # Makes the required connection matrix for the solver
        self.import_matrix()

        # Interpolation table:
        self.interpolation_grids()

        # Starts the solver.
        self.start_solver()

        stats = pstats.Stats(profiler).sort_stats('cumulative')
        stats.print_stats(10)

        # Postprocessing:
        self.post_graph()
        self.post()

    def load_yaml(self):
        """ Loads the provided input yaml file.
        """

        # # Reads the contents of the yaml file
        # with open(self.geometry_yaml_path) as file:
        #     contents = yaml.load(file, Loader=yaml.FullLoader)
        # # Sets the values in the class
        # for i in contents:
        #     setattr(self, i, contents[i])
        #
        # # Reads the contents of the yaml file
        # with open(self.simulation_yaml_path) as file:
        #     contents = yaml.load(file, Loader=yaml.FullLoader)
        # # Sets the values in the class
        # for i in contents:
        #     setattr(self, i, contents[i])
        #
        # # Reads the contents of the yaml file
        # with open(self.solver_settings_path) as file:
        #     contents = yaml.load(file, Loader=yaml.FullLoader)
        # # Sets the values in the class
        # for i in contents:
        #     setattr(self, i, contents[i])

        # # Reads the contents of the yaml file
        # with open(self.circuit_settings_path) as file:
        #     contents = yaml.load(file, Loader=yaml.FullLoader)
        # # Sets the values in the class
        # for i in contents:
        #     setattr(self, i, contents[i])

        # Reads the contents of the input yaml file
        with open(self.input_yaml_path) as file:
            contents = yaml.load(file, Loader=yaml.FullLoader)

        if 'path_dict' in contents.keys() and self.path_dict is None:
            setattr(self, 'path_dict', contents['path_dict'])
            contents.pop('path_dict', None)  # removing the key from the dictionary as we don't want to overwrite a custom path_dict.

        if self.path_dict is None:
            raise Exception("No path dictionary (path_dict) provided.")

        if 'geometry_yaml_path' not in self.path_dict.keys():
            raise Exception("No geometry path provided in the path dictionary.")

        if 'solver_settings_path' not in self.path_dict.keys():
            raise Exception("No solver settings path provided in the path dictionary.")

        if 'specific_output_path' not in self.path_dict.keys():
            raise Exception("No base path provided in the path dictionary.")

        for i in self.path_dict.keys():
            setattr(self, i, self.path_dict[i])

        if self.specific_output_path is None:
            raise Exception("No path to the simulation folder is provided.")

        # Checks if the geometry, and solver settings paths are provided, otherwise they will be loaded from the same folder as the input_yaml file.
        if self.geometry_yaml_path is None:
            self.geometry_yaml_path = Path(self.specific_output_path) / 'geometry.yaml'

        if self.solver_settings_path is None:
            self.solver_settings_path = Path(self.specific_output_path) / 'solver_settings.yaml'

        with open(self.geometry_yaml_path) as file:
            contents_geom = yaml.load(file, Loader=yaml.FullLoader)
        # Sets the values in the class
        for i in contents_geom:
            setattr(self, i, contents_geom[i])

        # Reads the contents of the yaml file
        with open(self.solver_settings_path) as file:
            contents_solv = yaml.load(file, Loader=yaml.FullLoader)
        # Sets the values in the class
        for i in contents_solv:
            setattr(self, i, contents_solv[i])

        # Sets the values in the class
        for i in contents:
            setattr(self, i, contents[i])

    def prepare_files(self):
        """ Prepares the result and logfiles.
                1) Makes the output folder
                3) Makes the output files
        """

        # Creates the general output folder if it doesn't exist.
        self.results_path = self.specific_output_path + os.sep + "results_light"

        try:
            os.mkdir(self.results_path)
        except OSError:
            print("Creation of the directory %s failed" % self.results_path)
        else:
            print("Successfully created the directory %s " % self.results_path)

        # Creates the specific simulation output folder.
        if self.sim_name == '':
            self.sim_name = datetime.now().strftime("%Y-%m-%d")

        self.sim_output = self.uniquify(self.results_path + os.sep + 'result' + ' - ' + self.sim_name)

        try:
            os.mkdir(self.sim_output)
        except OSError:
            print("Creation of the directory %s failed" % self.sim_output)
        else:
            print("Successfully created the directory %s " % self.sim_output)

        # Loading fieldmaps
        path = self.specific_output_path + os.sep + 'Bz_center.pkl'
        # Loading the full magnetic field matrices
        try:
            with open(path, 'rb') as fin:
                self.Bz_center = pickle.load(fin)
        except OSError:
            logging.info("Loading of %s has failed." % path)
        else:
            logging.info("Successfully loaded %s" % path)

        # Loading the full magnetic field matrices
        path = self.specific_output_path + os.sep + 'Br_center.pkl'
        try:
            with open(path, 'rb') as fin:
                self.Br_center = pickle.load(fin)
        except OSError:
            logging.info("Loading of %s has failed." % path)
        else:
            logging.info("Successfully loaded %s" % path)

        path = self.specific_output_path + os.sep + 'Bz_peak.pkl'
        # Loading the full magnetic field matrices
        try:
            with open(path, 'rb') as fin:
                self.Bz_peak = pickle.load(fin)
        except OSError:
            logging.info("Loading of %s has failed." % path)
        else:
            logging.info("Successfully loaded %s" % path)

        # Loading the full magnetic field matrices
        path = self.specific_output_path + os.sep + 'Br_peak.pkl'
        try:
            with open(path, 'rb') as fin:
                self.Br_peak = pickle.load(fin)
        except OSError:
            logging.info("Loading of %s has failed." % path)
        else:
            logging.info("Successfully loaded %s" % path)

        # Loading power dissipation matrix
        try:
            with open(self.power_dis_matrix_path, 'rb') as fin:
                self.power_dis_matrix = pickle.load(fin)
        except OSError:
            logging.info("Loading of %s has failed." % self.power_dis_matrix_path)
        else:
            logging.info("Successfully loaded %s" % self.power_dis_matrix_path)

        # Creating the files:
        self.file_summary = self.sim_output + os.sep + 'summary.csv'
        self.file_current = self.sim_output + os.sep + 'current.csv'
        self.file_current_R = self.sim_output + os.sep + 'current_resistor.csv'
        self.file_I_density =  self.sim_output + os.sep + 'I_density.csv'

        columns_1 = ['time'] + ['p' + str(i) for i in range(sum([len(i) for i in self.T0]))]
        columns_2 = ['time'] + ['p' + str(i) for i in range(len(self.Index_inductive_connections))]
        columns_3 = ['time'] + ['Stored Energy'] + ['I_' + str(i) for i in self.unique_circuits] + ['Bz Center'] + ['P radial'] + ['E dissipated'] + ['V' + str(i) for i in self.unique_circuits] + ['Solver time']

        try:
            with open(self.file_summary, 'w+') as f:
                f.write(self.delimiter.join(columns_3) + '\n')
        except OSError:
            logging.info("Creation of %s has failed." % self.file_summary)
        else:
            logging.info("Successfully created %s" % self.file_summary)

        try:
            with open(self.file_current, 'w+') as f:
                f.write(self.delimiter.join(columns_2) + '\n')
        except OSError:
            logging.info("Creation of %s has failed." % self.file_current)
        else:
            logging.info("Successfully created %s" % self.file_current)

        try:
            with open(self.file_current_R, 'w+') as f:
                f.write(self.delimiter.join(columns_1) + '\n')
        except OSError:
            logging.info("Creation of %s has failed." % self.file_current_R)
        else:
            logging.info("Successfully created %s" % self.file_current_R)

        try:
            with open(self.file_I_density, 'w+') as f:
                f.write(self.delimiter.join(columns_2) + '\n')
        except OSError:
            logging.info("Creation of %s has failed." % self.file_current_density)
        else:
            logging.info("Successfully created %s" % self.file_I_density)

        if self.export_geometry:
            shutil.copy(self.geometry_yaml_path, self.sim_output + os.sep + 'geometry.yaml')


    def create_ramp_rate_interpolation(self):
        """
        Creates interpolation function for the ramp rate.
        :return:
        """

        for i in self.unique_circuits:
            self.circuit_interp_dIdt[i] = interp1d(self.circuit_data[i]['t'], self.circuit_data[i]['dIdt'], bounds_error=False, fill_value=(self.circuit_data[i]['dIdt'][0], self.circuit_data[i]['dIdt'][-1]))

    def import_matrix(self):
        # Imports the required electrical matrices
        # These matrices are made by the geometry builder.

        if self.insulated:
            path = self.electrical_matrix_insulated_path
            try:
                with open(path, 'rb') as fin:
                    print(path)
                    self.electrical_matrix = pickle.load(fin)
            except OSError:
                logging.info("Loading of %s has failed." % path)
            else:
                logging.info("Successfully loaded %s" % path)

            path = self.mutual_inductance_matrix_insulated_path
            try:
                with open(path, 'rb') as fin:
                    print(path)
                    self.Mutual = pickle.load(fin)
            except OSError:
                logging.info("Loading of %s has failed." % path)
            else:
                logging.info("Successfully loaded %s" % path)

            self.inv = np.linalg.inv(self.electrical_matrix)  # Total inversed matrix (only the current part of it is needed in the ODE solver)
            self.inv_1 = self.inv[0: self.Nr_of_Connections_insulated, :]  # The current part
            self.inv_2 = self.inv[self.Nr_of_Connections_insulated:, :]  # The voltage part
            self.matrix_size = self.inv.shape[0]

        else:
            # Loading the electrical matrix
            path = self.electrical_matrix_path
            try:
                with open(path, 'rb') as fin:
                    print(path)
                    self.electrical_matrix = pickle.load(fin)
            except OSError:
                logging.info("Loading of %s has failed." % path)
            else:
                logging.info("Successfully loaded %s" % path)

            path = self.mutual_inductance_matrix_path
            try:
                with open(path, 'rb') as fin:
                    print(path)
                    self.Mutual = pickle.load(fin)
            except OSError:
                logging.info("Loading of %s has failed." % path)
            else:
                logging.info("Successfully loaded %s" % path)

            self.inv = np.linalg.inv(self.electrical_matrix)  # Total inversed matrix (only the current part of it is needed in the ODE solver)
            self.inv_1 = self.inv[0 : self.Nr_of_Connections, :]  # The current part
            self.inv_2 = self.inv[self.Nr_of_Connections:, :]  # The voltage part
            self.matrix_size = self.inv.shape[0]

    def interpolation_grids(self):
        # Makes two interpolation grids
        # grid_1 = Ic(T, B, deg), using the material function for the Ic
        # grid_2 = x(I/Ic, T, B)
        # Conductor_SC_properties= {'SC1': {'HTS_fit_name': 'CFUN_HTS_JcFit_Fujikura_v1', 'n-value': 15, 'Ic': 1500.0, 'Tref': 4.5, 'Bref': 40, 'Angle': 0.0}}
        # Conductor_SC_properties= {'SC1': {'HTS_fit_name': 'CFUN_HTS_JcFit_Fujikura_v1', 'n-value': 30, 'Ic': 91.2, 'Tref': 5.0, 'Bref': 5.0, 'Angle': 0.0, 'parallel': 6}}
        Conductor_SC_properties= {'SC1': {'HTS_fit_name': 'CFUN_HTS_JcFit_Fujikura_v1', 'n-value': 30, 'Ic': 342.2, 'Tref': 5.0, 'Bref': 5.0, 'Angle': 0.0, 'parallel': 6}}
        self.HTS_fit_name = 'CFUN_HTS_JcFit_Fujikura_v1'

        # Creating the normal and SC dictionaries
        self.SC_Ic_Fit = {}
        self.Normal_rho_FIT = {}
        self.SC_Cur_Sharing = {}
        self.SC_Ic_Scaling = {}

        # Creating a list of all normal conducting properties
        all_mats = list(np.concatenate([[j for j in self.Conductor_conductivity_materials[i]] for i in self.Conductor_conductivity_materials.keys()])) + [self.Conductor_resistor[i] for i in self.Conductor_resistor.keys()]
        all_mats = list(set(all_mats))
        print(all_mats)

        # Limits
        T = np.linspace(2.0, 100.0, 100)  # Temperature range
        T_rho = np.geomspace(2.0, 600.0, 100)  # Temperature range
        B = np.geomspace(0.1, 100.0, 100)  # Magnetic field range, geomspace, so 100 T is fine
        Degr = np.linspace(0.0, 360, 360)  # Angle
        T_torch = torch.tensor(T, dtype=torch.float64, device='cuda')
        T_rho_torch = torch.tensor(T_rho, dtype=torch.float64, device='cuda')
        B_torch = torch.tensor(B, dtype=torch.float64, device='cuda')
        Degr_torch = torch.tensor(Degr, dtype=torch.float64, device='cuda')
        IIc = np.geomspace(1e-9, 1000, 100)
        U = np.geomspace(1e-9, 2, 100)
        IIc_torch = torch.tensor(IIc, dtype=torch.float64, device='cuda')
        U_torch = torch.tensor(U, dtype=torch.float64, device='cuda')

        # Checks the fitname and turns angle dependency on or off.
        if self.HTS_fit_name in ['CFUN_HTS_JcFit_Fujikura_v1']:
            self.deg_dep = True
        else:
            self.deg_dep = False

        # Calculates the scaling factor: F_scaling = Ic_ref / Ic(T, B, deg)
        for i in Conductor_SC_properties.keys():
            if self.deg_dep:
                Tref = Conductor_SC_properties['SC1']['Tref']
                Bref = Conductor_SC_properties['SC1']['Bref']
                Deg = Conductor_SC_properties['SC1']['Angle']
                Ic = STEAM_materials(self.HTS_fit_name, 3, 1).evaluate(np.array([Tref, Bref, Deg]))
                self.SC_Ic_Scaling[i] = Conductor_SC_properties['SC1']['Ic'] / Ic
            else:
                print('Not implemented yet')

            # Making the Ic interpolation
            input_list = np.array([i for i in itertools.product(T, B, Degr)])
            temp = np.vstack((input_list[:, 0], input_list[:, 1], input_list[:, 2]))
            if self.deg_dep:
                Ic = STEAM_materials(self.HTS_fit_name, temp.shape[0], temp.shape[1]).evaluate(temp) * self.SC_Ic_Scaling[i]
                Ic_torch    = torch.tensor(Ic.reshape((len(T), len(B), len(Degr))), dtype=torch.float64, device='cuda')
                # print(len(T_torch), len(B_torch), len(Degr_torch), len(Ic_torch))
                self.SC_Ic_Fit[i] = RegularGridInterpolator((T_torch, B_torch, Degr_torch), Ic_torch)
            else:
                print('Not yet implemented')

        # Making the resistivity interpolation function
        for i in all_mats:
            if i == 'Cu30':
                rrr = [50]
                input_list = np.array([i for i in itertools.product(T_rho, B, rrr)])
                temp = np.vstack((input_list[:, 0], input_list[:, 1], input_list[:, 2]))
                R = STEAM_materials('CFUN_rhoCu_NIST_v2', temp.shape[0], temp.shape[1]).evaluate(temp)
                R_torch = torch.tensor(R.reshape((len(T), len(B))), dtype=torch.float64, device='cuda')
                self.Normal_rho_FIT[i] = RegularGridInterpolator((T_rho_torch, B_torch), R_torch)
            else:
                input_list = np.array([i for i in itertools.product(T_rho, B)])
                temp = np.vstack((input_list[:, 0], input_list[:, 1]))
                print(temp)
                R = STEAM_materials('CFUN_rhoHast_v1', 1, len(temp[0,:])).evaluate(temp[0,:])
                R_torch = torch.tensor(R.reshape((len(T), len(B))), dtype=torch.float64, device='cuda')
                self.Normal_rho_FIT[i] = RegularGridInterpolator((T_rho_torch, B_torch), R_torch)

        # Making the x interpolation function
        for i in Conductor_SC_properties.keys():
            input_list = np.array([k for k in itertools.product(IIc, U)])
            x = np.zeros((len(input_list), 1))
            for j in range(len(input_list)):
                # print((input_list[j][0], input_list[j][1], Conductor_SC_properties[i]['n-value']))
                x[j] = fsolve(self.powerlaw, 1.0, args=(input_list[j][0], input_list[j][1], Conductor_SC_properties[i]['n-value']), maxfev=1000000)
            x_torch = torch.tensor(x.reshape((len(IIc), len(U))), dtype=torch.float64, device='cuda')
            self.SC_Cur_Sharing[i] = RegularGridInterpolator((IIc_torch, U_torch), x_torch)

    @staticmethod
    def powerlaw(x, IIc, U, n=21, E0=1e-4):
        return E0 * IIc ** n * x ** n - U * (1 - x)

    def PID(self, Power, dt=0):
        # Power is target power.
        # Value of offset - when the error is equal zero
        offset = 0.0

        # PID calculations
        e = self.setpoint - Power

        P = self.Kp*e
        self.integral += self.Ki*e*dt
        # D = self.Kd*(e - e_prev)/dt

        # calculate manipulated variable - MV
        rr = offset + P + self.integral # + D

        # update stored data for next iteration
        # e_prev = e

        if abs(rr) > self.max_rr and rr > 0: rr = np.sign(rr) * self.max_rr
        if abs(rr) > abs(self.min_rr) and rr < 0: rr = self.min_rr
        return rr

    def start_solver(self):
        # Pre-processing some things:
        self.cross_sections = np.array(self.cross_sections)
        self.turns_per_element = np.array(self.turns_per_element)
        self.lengths = np.array(self.lengths)

        # if self.steam_mat_Ic:
        #     Tref = self.SC_properties['Tref']
        #     Bref = self.SC_properties['Bref']
        #     Deg = self.SC_properties['angle']
        #     Ic = STEAM_materials(self.SC_properties['fit-name'], 3, 1).evaluate(np.array([Tref, Bref, Deg]))
        #     self.SC_Ic_Scaling = self.SC_properties['Ic'] / Ic

        # Initial Conditions
        t = 0.0
        dt = self.dt
        dt = 1.0
        if self.insulated:
            x0 = np.zeros(self.Nr_of_Connections_insulated)
            self.lengths = self.lengths[:self.Nr_of_Connections_insulated]
        else:
            x0 = np.zeros(self.Nr_of_Connections)

        # self.radial_resistance = np.array(self.radial_resistance[0] + self.radial_resistance[1])
        radial_resistance_temp = []
        [radial_resistance_temp.extend(i) for i in self.radial_resistance]
        self.radial_resistance = np.array(radial_resistance_temp)
        if self.R_multiplier is not None:
            if type(self.R_multiplier) in [int, float, np.float64]:
                self.radial_resistance = self.radial_resistance * self.R_multiplier

        if self.insulated:
            # Increases the radial resistances to 1 kOhm
            print("The Insulation is set to True, increasing the turn-to-turn resistance to 1 kOhm.")
            self.radial_resistance = np.where(self.radial_resistance > 0, 1e3, self.radial_resistance)

        self.superconductor = np.where(np.abs(self.radial_resistance) > 0, 0.0, 1.0)

        print('*** Starting ODE Solver ***')
        profiler.enable()
        start = time.time_ns()*1e-9
        if self.solver == 'vode':
            self.results = scipy.integrate.ode(self.solve_this).set_integrator('vode', method='bdf', rtol=self.rtol, atol=self.atol, min_step=self.min_step, max_step=self.max_step, nsteps=int(1e7), order=5)
        self.results.set_initial_value(x0, t)

        # The solver while loop.
        self.t_array.append(t)
        self.I_array.append(x0)

        # I_density
        self.I_density_array.append(x0[self.Index_inductive_connections] / self.cross_sections[self.Index_inductive_connections] * self.turns_per_element[self.Index_inductive_connections] * 1e-6)
        self.solver_t_array.append(time.time_ns() * 1e-9 - start)
        if self.fast:
            while self.results.successful() and self.results.t < self.t0[1]:  # and ((E > self.stop_criterion_E) & self.stop_on_E):
                self.results.integrate(self.results.t + dt)
                self.t_array.append(self.results.t)
                self.I_array.append(self.results.y[:self.Nr_of_Connections])
                self.solver_t_array.append(time.time_ns()*1e-9-start)
                self.I_density_array.append(self.results.y[:self.Nr_of_Connections][self.Index_inductive_connections] / self.cross_sections[self.Index_inductive_connections] * self.turns_per_element[self.Index_inductive_connections] * 1e-6)

        else:
            while self.results.successful() and self.results.t < self.t0[1]:  # and ((E > self.stop_criterion_E) & self.stop_on_E):
                self.results.integrate(self.results.t + dt)
                if np.mod(self.save_n, self.savemod) == 0:
                    self.t_array.append(self.results.t)
                    self.I_array.append(self.results.y[:self.Nr_of_Connections])
                    self.solver_t_array.append(time.time_ns() * 1e-9 - start)
                    self.I_density_array.append(self.results.y[:self.Nr_of_Connections][self.Index_inductive_connections] / self.cross_sections[self.Index_inductive_connections] * self.turns_per_element[self.Index_inductive_connections] * 1e-6)

                ### PID TEST ###
                # Go to nominal current as fast as possible while staying below x amount of dissipated heat.
                # Ramp rate limited to dI_lim
                I_circuit, Power = self.solve_this(self.results.t, self.results.y[:self.Nr_of_Connections], False, True)
                self.setpoint = 200 if (abs(I_circuit[0]) < 400.0 and self.setpoint == 200) else 100
                self.rr = self.PID(Power, dt) if abs(I_circuit[0]) < self.nominal_current else 0.0
                print(self.results.t, self.rr, I_circuit, Power)
                self.save_n += 1

        print('Starting post processing.')
        for i in range(len(self.t_array)):
            print(str(i) + ' out of ' + str(len(self.t_array)))
            self.solve_this(self.t_array[i], self.I_array[i], True)
            self.Bz_center_array.append(np.dot(self.Bz_center, np.array(self.I_array[i])[self.Index_inductive_connections])[0])
            self.stored_energy.append(np.sum(0.5 * self.Mutual * np.outer(np.array(self.I_array[i]), np.array(self.I_array[i]))))
        self.E_dis = integrate.cumtrapz(self.Power, self.t_array, initial=0)

        print('Starting export.')

        self.save_results()

        # End of simulation
        end = time.time_ns()*1e-9
        now = datetime.now().strftime("%H:%M:%S")
        print(str(now) + ' || Elapsed time: ' + str((end - start)) + ' seconds.')
        profiler.disable()
        self.I_array = np.array(self.I_array)

        fig = plt.figure()
        ax = fig.add_subplot()
        for i in (range(self.I_array.shape[1])):
            ls = '-' if i < 24 else '--'
            ax.plot(self.t_array, self.I_array[:,i], ls=ls)

        ax.set_xlabel('Time, s')
        ax.set_ylabel('Current, A')
        ax.set_title('Stepsize: ' +str(dt) + 's, elapsed time: ' + str(np.round(end - start,3)) + 's')
        plt.show()

    # def solve_this_torch(self, t, x, output2 = False, output3=False):
    #     """ This is the function that is solved by the ode solver.
    #         It basically does this: dI/dt = f(t, I) = inv(M) * b
    #         With b containing the resistive voltages in the system and any external dI/dt's.
    #     """
    #     x_torch = torch.tensor(x, dtype=torch.float64, device='cuda')
    #     # I_ind =
    #     # I_res =
    #     b = torch.zeros((self.inv.shape[0],)) #TODO move up and make it like this b = self.b0.clone().detach()
    #
    #     # Magnetic Field Calculations
    #     # angle = arccos(Br / sqrt(Bz^2 + Br^2)), 0 deg = normal to the tape, 90 deg = perpendicular to the tape
    #     Bz = torch.matmul(self.Bz_Conductor, I_ind)
    #     Br = torch.matmul(self.Br_Conductor, I_ind)
    #     B_tot = torch.sqrt(Bz^2 + Br^2)
    #     angle = torch.arccos(Br / B_tot)
    #
    #     # Calculating the current sharing:
    #     # 1: calculate the Ic per SC
    #
    #     # 2: calculate the resistivity of the stabilizer materials
    #
    #     # 3: calculate the current sharing coefficient
    #
    #     # 4: calculate the voltage
    #
    #     # 5: calculate the power
    #
    #     # The dot product of the inversed matrix with the b-array.
    #     dIdt = torch.matmut(self.inv_1, b) #TODO make the inverse a torch tensor
    #     return dIdt.cpu().numpy()

    def solve_this(self, t, x, output2 = False, output3=False):
        """ This is the function that is solved by the ode solver.
            It basically does this: dI/dt = f(t, I) = inv(M) * b
            With b containing the resistive voltages in the system and any external dI/dt's.
        """
        # Separating the current from the temperature:
        # I = np.array(x[0:self.Nr_of_Connections])  # All Currents

        # Making the b-array filled with zeros.
        b = np.zeros((self.inv.shape[0],))
        # Ramp
        # b[self.circuits_list[1]] += self.rr #self.circuit_interp_dIdt[1](t)
        b[self.circuits_list[1]] += self.circuit_interp_dIdt[1](t)
        if not self.insulated:
            b[self.b_V_array_index] += -(self.radial_resistance * x)  # Voltage due to radial resistance and current

        if self.include_Ic and self.insulated:
            # TODO brickolage to generate sim results for gianluca
            if self.steam_mat_Ic:
                Bz = np.dot(self.Bz_peak, x)
                Br = np.dot(self.Br_peak, x)
                B_tot = np.sqrt(Bz ** 2 + Br ** 2)
                angle = np.nan_to_num(np.abs(np.arccos(Br / B_tot)) / np.pi * 180)  #* 0.0
                temp = np.vstack((4.5*np.ones_like(x), B_tot, angle))
                Ic = STEAM_materials(self.SC_properties['fit-name'], temp.shape[0], temp.shape[1]).evaluate(temp) * self.SC_Ic_Scaling['SC1']

                # From the Torch solver:
                R = np.zeros_like(x)
                T = torch.ones(x.shape, dtype=torch.float64, device='cuda') * 4.5
                B_torch = torch.tensor(np.clip(B_tot, a_min=0.11, a_max=100), dtype=torch.float64, device='cuda')
                # print(self.Conductor_materials)
                # print(self.Conductor_materials_cross_section)
                for i in range(len(self.Conductor_materials['SC1'])): # TODO hardcoded
                    # print(self.Conductor_materials['SC1'][i], self.Conductor_materials_cross_section['SC1'][i])
                    # print(T, B_tot)
                    R += self.Normal_rho_FIT[self.Conductor_materials['SC1'][i]]((T, B_torch)).cpu().numpy() / self.Conductor_materials_cross_section['SC1'][i]
                R = 1/R
                IIc = np.abs(x/Ic)
                # print(IIc)
                U = np.abs(R * x)
                IIc = torch.tensor(IIc, dtype=torch.float64, device='cuda')
                U = torch.tensor(U, dtype=torch.float64, device='cuda')
                Rt = torch.tensor(R, dtype=torch.float64, device='cuda')
                # print(self.SC_Cur_Sharing['SC1']((IIc, U)))
                R_Conductor = torch.abs((Rt * (1 - self.SC_Cur_Sharing['SC1']((IIc, U))))).cpu().numpy() # TODO hardcoded key.
                # print(t, np.max(Ic), np.min(Ic), np.max(angle), np.min(angle))
                # v = 1e-4 * np.clip(abs(x / np.clip(Ic, a_min=0.1, a_max=100000000)), a_min=0.0, a_max=1000) ** self.SC_properties['n-value'] * np.sign(x)
                v = R_Conductor * x * self.lengths
            else:
                v = 1e-4 * (x / 200.0) ** 30
            # v_st = (3e-4/50)*abs(x)
            # v_st = 1e-6*abs(x)
            # v = np.where(abs(v) > v_st, v_st*np.sign(x), v)
            # b[-len(x):] += -v * self.lengths  # Voltage due to the SC current
            b[-len(x):] -= v
            # print('Time (s): ',t)

        # The dot product of the inversed matrix with the b-array.
        dIdt = np.dot(self.inv_1, b)
        # Output 1: only for the solver.
        if output2:
            V = np.dot(self.inv_2, b)
            if self.insulated:
                self.I_circuit.append([0.0 for i in self.unique_circuits])
                self.V_circuit.append([np.sum(V[self.circuits_list[i]]) for i in self.unique_circuits])
                self.Power.append(0.0)
            else:
                self.I_circuit.append([(np.dot(self.power_dis_matrix, x))[self.circuits_list[i]][0] for i in self.unique_circuits])
                self.V_circuit.append([np.sum(V[self.circuits_list[i]]) for i in self.unique_circuits])
                self.Power.append(np.sum(self.radial_resistance * x ** 2))
        elif output3:
            I_circuit = [(np.dot(self.power_dis_matrix, x))[self.circuits_list[i]][0] for i in self.unique_circuits]
            Power = np.sum(self.radial_resistance * x ** 2)
            return I_circuit, Power
        else:
            return dIdt

    def save_results(self):
        saverange = np.arange(len(self.t_array))

        try:
            with open(self.file_summary, 'a+') as f:
                for j in saverange:
                    # print('E', self.stored_energy[j])
                    # print('Icircuit', [i for i in self.I_circuit[j]])
                    # print('Bz', self.Bz_center_array[j])
                    # print('P', self.Power[j])
                    # print('Vcircuit', [i for i in self.V_circuit[j]])
                    # print('Solver time', self.solver_t_array[j])
                    data = [str(self.t_array[j])] + [self.output_format(self.stored_energy[j])] + [self.output_format(i) for i in self.I_circuit[j]] + [self.output_format(self.Bz_center_array[j])] + [self.output_format(self.Power[j])] + [self.output_format(self.E_dis[j])] + [self.output_format(i) for i in self.V_circuit[j]] + [self.output_format(self.solver_t_array[j])]
                    f.write(self.delimiter.join(data) + '\n')
        except OSError:
            logging.info("Writing on %s has failed." % self.file_summary)

        try:
            with open(self.file_current, 'a+') as f:
                for j in saverange:
                    data = [str(self.t_array[j])] + [self.output_format(i) for i in np.array(self.I_array[j])]
                    f.write(self.delimiter.join(data) + '\n')
        except OSError:
            logging.info("Writing on %s has failed." % self.file_current)

        # try:
        #     with open(self.file_current_R, 'a+') as f:
        #         for j in saverange:
        #             data = [str(self.t_array[j])] + [self.output_format(i) for i in np.array(self.I_array[j])]
        #             f.write(self.delimiter.join(data) + '\n')
        # except OSError:
        #     logging.info("Writing on %s has failed." % self.file_current_R)

        try:
            with open(self.file_I_density, 'a+') as f:
                for j in saverange:
                    data = [str(self.t_array[j])] + [self.output_format(i) for i in np.array(self.I_density_array[j])]
                    f.write(self.delimiter.join(data) + '\n')
        except OSError:
            logging.info("Writing on %s has failed." % self.file_I_density)

    @staticmethod
    def output_format(nr, form = '{:.5E}'):
        # 3 decimals = .3E, 4 decimals = .4E etc., E = scientific. F = float.
        return form.format(nr)

    def post_graph(self):
        fig = plt.figure(figsize=(16, 7))

        ax1 = fig.add_subplot(241)
        ax2 = fig.add_subplot(242)
        ax3 = fig.add_subplot(243)
        ax4 = fig.add_subplot(244)
        ax5 = fig.add_subplot(245)
        ax6 = fig.add_subplot(246)
        ax7 = fig.add_subplot(247)

        ax1.plot(self.t_array, self.stored_energy)
        ax2.plot(self.t_array, self.Bz_center_array)
        ax3.plot(self.t_array, self.Power)
        [ax4.plot(self.t_array, np.array(self.I_circuit)[:, i]) for i in range(len(self.I_circuit[0]))]
        [ax5.plot(self.t_array, np.array(self.V_circuit)[:, i]) for i in range(len(self.V_circuit[0]))]
        ax6.plot(self.t_array, self.solver_t_array)
        ax7.plot(self.t_array, self.E_dis)

        ax1.set_xlabel('Time, s')
        ax1.set_ylabel('Stored Energy, J')
        ax2.set_xlabel('Time, s')
        ax2.set_ylabel('Bz Center, T')
        ax3.set_xlabel('Time, s')
        ax3.set_ylabel('Radial Loss, W')
        ax4.set_xlabel('Time, s')
        ax4.set_ylabel('Current, A')
        ax5.set_xlabel('Time, s')
        ax5.set_ylabel('Voltage, V')
        ax6.set_xlabel('Time, s')
        ax6.set_ylabel('Solving Time, s')
        ax7.set_xlabel('Time, s')
        ax7.set_ylabel('Dissipated Energy, J')

        fig.tight_layout()

        fig.savefig(self.sim_output + os.sep + 'results.png', dpi=300)
        plt.show()
        plt.close(fig)

    def post(self):
        # Loading fieldmaps
        path = self.path + os.sep + 'Bz_full.pkl'
        try:
            with open(path, 'rb') as fin:
                self.Bz_full = pickle.load(fin)
        except OSError:
            print("Loading of %s has failed." % (path))
        else:
            print("Successfully loaded %s" % (path))
        path = self.path + os.sep + 'Br_full.pkl'
        try:
            with open(path, 'rb') as fin:
                self.Br_full = pickle.load(fin)
        except OSError:
            print("Loading of %s has failed." % (path))
        else:
            print("Successfully loaded %s" % (path))

        x = 10
        y = 10
        xmin, xmax, ymin, ymax, nx, ny = self.grid_information

        Bz_array = []
        Br_array = []

        for i in self.I_array:
            Bz = np.dot(self.Bz_full, np.array(i)[self.Index_inductive_connections]).reshape((ny, nx))[x,y]
            Br = np.dot(self.Br_full, np.array(i)[self.Index_inductive_connections]).reshape((ny, nx))[x,y]
            Bz_array.append(Bz)
            Br_array.append(Br)

        fig = plt.figure()
        ax = fig.add_subplot()

        ax.plot(self.t_array, Bz_array, label='Bz')
        ax.plot(self.t_array, Br_array, label='Br')
        ax.legend()

        ax.set_xlabel('Time, s')
        ax.set_ylabel('Magnetic Field, T')
        ax.set_title('Magnetic Field at position: (' + str(np.round(np.linspace(xmin, xmax, nx)[x],3)) + ',' + str(np.round(np.linspace(ymin, ymax, ny)[y], 3)) + ')')
        fig.tight_layout()
        plt.show()

    def post2(self):

        Bz_array = []
        Br_array = []

        for i in self.I_array:
            Bz = np.dot(self.Bz_center, np.array(i)[self.Index_inductive_connections])
            Br = np.dot(self.Br_center, np.array(i)[self.Index_inductive_connections])
            Bz_array.append(Bz)
            Br_array.append(Br)

        fig = plt.figure()
        ax = fig.add_subplot()

        ax.plot(self.t_array, Bz_array, label='Bz')
        ax.plot(self.t_array, Br_array, label='Br')
        ax.legend()

        ax.set_xlabel('Time, s')
        ax.set_ylabel('Magnetic Field, T')
        ax.set_title('Magnetic Field at position: (0.0, 0.0)')
        fig.tight_layout()
        plt.show()

    @staticmethod
    def uniquify(path_original):
        # Creates a unique foldername for the simulation
        counter = 1
        path = path_original + " - " + str(counter)

        while os.path.isdir(path):
            counter += 1
            path = path_original + " - " + str(counter)

        return path


if __name__ == "__main__":
    import cProfile, pstats
    profiler = cProfile.Profile()
    NI_sim = NI_solver_light(sys.argv[1])