#!/usr/bin/env python3

import numpy as np
import matplotlib.pyplot as plt
from matplotlib import rcParams
from matplotlib.ticker import MultipleLocator
from matplotlib.colors import LightSource
import matplotlib
import torch
import os
import yaml
import pandas as pd
from itertools import cycle
from steammaterials.STEAM_materials import STEAM_materials
from scipy import integrate  
from typing import List, Optional
from .materials import MaterialProperties

rcParams['axes.labelsize'] = 11
rcParams['xtick.labelsize'] = 11
rcParams['ytick.labelsize'] = 11
rcParams['legend.fontsize'] = 11
rcParams['font.family'] = 'DejaVu Sans'
rcParams['font.sans-serif'] = 'Arial'
rcParams['text.usetex'] = False
rcParams['mathtext.fontset'] = 'dejavusans'
rcParams['figure.autolayout'] = True
rcParams['lines.markersize'] = 5
rcParams['figure.figsize'] = (6, 5)
rcParams['axes.grid'] = True

class DebuggingGraphs:
    """
    A library containing debugging graph generation functions for the NI coil solver.
    These functions generate diagnostic plots for resistance, critical current, and current sharing.
    """
    
    def __init__(self, solver):
        """
        Initialize the debugging graphs with a reference to the solver instance.
        
        Args:
            solver: The NI_Coil_Solver instance
        """
        self.solver = solver
    
    def generate_radial_resistance_graph(self):
        """
        Generates a graph with all the radial resistances. Used for debugging purposes.
        :return:
        """

        # Graph after R calculation:
        fig = plt.figure(figsize=(12, 4))
        ax1 = fig.add_subplot(121)
        ax2 = fig.add_subplot(122)

        for i in self.solver.Index_initial_IR.keys():
            ax1.plot(self.solver.Index_initial_IR[i], self.solver.ref_resistivity[self.solver.Index_initial_IR[i]].cpu().numpy())
            ax2.plot(self.solver.Index_initial_IR[i], self.solver.ref_resistance[self.solver.Index_initial_IR[i]].cpu().numpy())

        ax1.set_xlabel('Position Indexed, -')
        ax1.set_ylabel('Resistivity, $\\Omega$m')
        ax2.set_xlabel('Position, -')
        ax2.set_ylabel('Resistance, $\\Omega$')

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

    def generate_ic_graph(self):
        """
        Generates the Ic graphs for all unique SC regions.
        :return:
        """

        B = np.linspace(0.1, 50, 1000)
        T = np.linspace(2.0, 95, 1000)
        T_initial = np.ones_like(B) * np.mean(self.solver.T0)
        Deg = [0.0, 90.0]

        fig_B = plt.figure()
        ax_B = fig_B.add_subplot(111)
        fig_T = plt.figure()
        ax_T = fig_T.add_subplot(111)

        ls_array = ['-', '--', '-.', ':']
        color_array = ['blue', 'orange', 'green', 'red', 'black', 'purple']
        colors = cycle(color_array)

        # Ic as a function of magnetic field.
        for i in self.solver.unique_regions:
            color = next(colors)
            if self.solver.Conductor_type[i] == 'Superconductor':
                ls = cycle(ls_array)
                for deg in Deg:
                    lss = next(ls)
                    fit_name = self.solver.Conductor_SC_properties[i]['fit-name']
                    temp = MaterialProperties.create_steam_material_input(
                        fit_name, 
                        Temperature=T_initial, 
                        Magnetic_Field=B, 
                        Field_Angle=np.ones_like(B)*deg
                    )
                    Ic = STEAM_materials(fit_name, temp.shape[0], temp.shape[1]).evaluate(temp) * self.solver.SC_Ic_Scaling[i].cpu().numpy()
                    ax_B.plot(B, Ic, ls=lss, color=color, label=str(i) + ', ' + str(deg) + ' Degrees')
        ax_B.legend()
        ax_B.set_yscale('log')
        ax_B.set_xlabel('Magnetic Field, T')
        ax_B.set_ylabel('Critical Current per Element, A')

        # Ic as a function of temperature.
        colors = cycle(color_array)
        B_initial = 5.0
        for i in self.solver.unique_regions:
            color = next(colors)
            if self.solver.Conductor_type[i] == 'Superconductor':
                ls = cycle(ls_array)
                for deg in Deg:
                    lss = next(ls)
                    fit_name = self.solver.Conductor_SC_properties[i]['fit-name']
                    temp = MaterialProperties.create_steam_material_input(
                        fit_name, 
                        Temperature=T, 
                        Magnetic_Field=np.ones_like(T) * B_initial, 
                        Field_Angle=np.ones_like(T) * deg
                    )
                    Ic = STEAM_materials(fit_name, temp.shape[0], temp.shape[1]).evaluate(temp) * self.solver.SC_Ic_Scaling[i].cpu().numpy()
                    ax_T.plot(T, Ic, ls=lss, color=color, label=str(i) + ', ' + str(deg) + ' Degrees, 5 T')

        ax_T.legend()
        bottom, top = ax_T.get_ylim()
        ax_T.set_ylim((1e-1, top)) # Setting the lower limit of the ax to 1e-1 A.
        ax_T.set_yscale('log')
        ax_T.set_xlabel('Temperature, K')
        ax_T.set_ylabel('Critical Current per Element, A')

        fig_B.savefig(self.solver.sim_output + os.sep + 'IcB.png', dpi=300)
        fig_T.savefig(self.solver.sim_output + os.sep + 'IcT.png', dpi=300)
        plt.close(fig_B)
        plt.close(fig_T)

    def generate_current_sharing_surface(self):
        """
        Generates a graph of the current sharing coefficient as a function of I, B and T for all the unique superconducting regions.
        :return:
        """

        B = np.linspace(0, 25.0, 150)
        T = np.linspace(20, 100.0, 150)
        Angle = 0.0
        I = 250.0 #TODO this is hardcoded, should not be and selected from input parameters, such as Ic.

        xv, yv = np.meshgrid(B, T)

        for i in self.solver.unique_regions:
            if self.solver.Conductor_type[i].lower() == 'superconductor':
                fig_x = plt.figure()
                ax_x = fig_x.add_subplot(111, projection='3d')

                # Getting the data
                fit_name = self.solver.Conductor_SC_properties[i]['fit-name']
                temp = MaterialProperties.create_steam_material_input(
                    fit_name,
                    Temperature=yv.ravel(),
                    Magnetic_Field=xv.ravel(),
                    Field_Angle=Angle * np.ones_like(xv.ravel())
                )
                Ic = STEAM_materials(fit_name, temp.shape[0], temp.shape[1]).evaluate(temp) * self.solver.SC_Ic_Scaling[i].cpu().numpy()
                Ic = np.clip(Ic, 1e-99, None)

                # Resistance of superconductor i
                R = np.zeros_like(xv.ravel())
                for l in range(len(self.solver.Conductor_materials[i])):
                    rho = self.solver.materials.resistivity_selector(self.solver.Conductor_materials[i][l], yv.ravel(), xv.ravel(), reshape=True)
                    R += 1 / (rho / self.solver.Conductor_materials_cross_section[i][l])
                R = (1 / R) * self.solver.Conductor_SC_properties[i]['parallel']

                # For the x interpolation function
                IIc = torch.tensor(np.abs(I / Ic), dtype=self.solver.torch_float_type, device=self.solver.torch_device)
                U = torch.tensor(np.abs(R * I), dtype=self.solver.torch_float_type, device=self.solver.torch_device)

                x = self.solver.SC_Cur_Sharing[i]((IIc, U)).cpu().numpy()

                # Use single color with grid lines for publication quality
                # rstride and cstride make grid lines sparser
                ax_x.plot_surface(xv, yv, x.reshape(xv.shape), color='lightsteelblue',
                                 rstride=5, cstride=5, linewidth=0.5, edgecolor='gray',
                                 antialiased=True, alpha=0.9, shade=False)

                ax_x.set_xlabel('Magnetic Field, T')
                ax_x.set_ylabel('Temperature, K')
                ax_x.set_zlabel('Current Sharing Coefficient, -')

                # Rotate camera by 90 degrees around z-axis
                ax_x.view_init(elev=30, azim=30)

                fig_x.savefig(self.solver.sim_output + os.sep + str(i) + '.png', dpi=300)
                plt.close(fig_x)

    def generate_ic_surface(self):
        """
        Generates a 3D surface plot of critical current (Ic) as a function of
        temperature and magnetic field at a fixed field angle of 0 degrees
        for all unique superconducting regions.

        This visualization helps understand how the critical current depends
        on both temperature and magnetic field simultaneously.
        :return:
        """

        B = np.linspace(0, 25.0, 150)
        T = np.linspace(20, 100.0, 150)
        Angle = 0.0  # Fixed field angle in degrees

        xv, yv = np.meshgrid(B, T)

        for i in self.solver.unique_regions:
            if self.solver.Conductor_type[i].lower() == 'superconductor':
                fig_ic = plt.figure()
                ax_ic = fig_ic.add_subplot(111, projection='3d')

                # Getting the Ic data
                fit_name = self.solver.Conductor_SC_properties[i]['fit-name']
                temp = MaterialProperties.create_steam_material_input(
                    fit_name,
                    Temperature=yv.ravel(),
                    Magnetic_Field=xv.ravel(),
                    Field_Angle=Angle * np.ones_like(xv.ravel())
                )
                Ic = STEAM_materials(fit_name, temp.shape[0], temp.shape[1]).evaluate(temp) * self.solver.SC_Ic_Scaling[i].cpu().numpy()
                Ic = np.clip(Ic, 1e-99, None)

                # Use single color with grid lines for publication quality
                # rstride and cstride make grid lines sparser
                ax_ic.plot_surface(xv, yv, Ic.reshape(xv.shape), color='lightsteelblue',
                                  rstride=5, cstride=5, linewidth=0.5, edgecolor='gray',
                                  antialiased=True, alpha=0.9, shade=False)

                ax_ic.set_xlabel('Magnetic Field, T')
                ax_ic.set_ylabel('Temperature, K')
                ax_ic.set_zlabel('Critical Current, A')
                ax_ic.set_title(f'Ic vs T and B (Angle={Angle}°)')

                # Rotate camera by 90 degrees around z-axis
                ax_ic.view_init(elev=30, azim=30)

                fig_ic.savefig(self.solver.sim_output + os.sep + str(i) + '_Ic_surface.png', dpi=300)
                plt.close(fig_ic)

    def generate_ic_vs_angle_graph(self):
        """
        Generates a graph of critical current as a function of field angle
        at fixed temperature (5 K) and magnetic field (5 T) for all unique
        superconducting regions.

        This is useful for visualizing the angular dependence of Ic, which
        can vary significantly in anisotropic superconductors like REBCO.
        :return:
        """
        # Fixed conditions for the angle sweep
        T_fixed = self.solver.Ic_T_fixed # K
        B_fixed = self.solver.Ic_B_fixed # T
        Angle = np.linspace(0.0, 180.0, 200)  # Degrees, 90 = parallel, 0 = perpendicular

        fig = plt.figure(figsize=(8, 6))
        ax = fig.add_subplot(111)

        # Color cycle for different regions
        color_array = ['blue', 'orange', 'green', 'red', 'black', 'purple']
        colors = cycle(color_array)

        # Plot Ic vs angle for each superconducting region
        for i in self.solver.unique_regions:
            if self.solver.Conductor_type[i] == 'Superconductor':
                color = next(colors)
                fit_name = self.solver.Conductor_SC_properties[i]['fit-name']

                # Check if angle mirroring is enabled for this region
                mirror_angle = self.solver.Conductor_SC_properties[i].get('mirror_angle', False)

                # Apply angle mirroring if enabled
                # Mirrors angles around 90 deg (e.g., 150->30, 40->140, 90->90)
                if mirror_angle:
                    Angle_effective = 180 - Angle
                else:
                    Angle_effective = Angle

                # Create input array with fixed T and B, varying angle
                temp = MaterialProperties.create_steam_material_input(
                    fit_name,
                    Temperature=np.ones_like(Angle) * T_fixed,
                    Magnetic_Field=np.ones_like(Angle) * B_fixed,
                    Field_Angle=Angle_effective
                )

                # Evaluate critical current using STEAM materials model
                Ic = STEAM_materials(fit_name, temp.shape[0], temp.shape[1]).evaluate(temp) * self.solver.SC_Ic_Scaling[i].cpu().numpy()

                # Create label with mirror indication
                label = f'Region {i}'
                if mirror_angle:
                    label += ' (mirrored)'

                # Plot with region label
                ax.plot(
                    Angle,
                    Ic,
                    color=color,
                    linewidth=2,
                    label=label
                )

        # Configure plot appearance
        ax.set_xlabel('Field Angle, deg')
        ax.set_ylabel('Critical Current, A')
        ax.set_title(f'Critical Current vs Field Angle at {T_fixed} K and {B_fixed} T')
        ax.legend(loc='best')
        ax.grid(True, alpha=0.3)

        # Save the figure
        fig.savefig(self.solver.sim_output + os.sep + 'IcAngle.png', dpi=300)
        plt.close(fig)

    def generate_time_constant_vs_temperature(self):
        """
        Generates graphs of magnet time constant (tau = L/R) as a function
        of temperature at constant magnetic field for each circuit.

        For No-Insulation coils, the time constant is determined by the
        radial (turn-to-turn) resistance. Temperature dependence comes from
        material resistivity changes with temperature.
        :return:
        """
        # Temperature range for the sweep
        T = np.linspace(2.0, 80.0, 200)  # K
        B_fixed = 0.2  # Fixed magnetic field in T

        # Create figure for each circuit
        for circuit_id in self.solver.unique_circuits:
            if circuit_id not in self.solver.inductance_dict:
                continue
            if circuit_id not in self.solver.index_dict['circuits']:
                continue

            # Get circuit inductance (constant)
            L_circuit = self.solver.inductance_dict[circuit_id]

            # Get radial connection indices for this circuit
            radial_indices = (
                self.solver.index_dict['circuits'][circuit_id]['IR']
            )

            # Skip if no radial elements
            if len(radial_indices) == 0:
                continue

            # Calculate total circuit resistance as function of temperature
            R_circuit = np.zeros_like(T)

            for idx, temp in enumerate(T):
                R_total = 0.0

                # Sum radial resistance of all elements in this circuit
                for radial_idx in radial_indices:
                    # Find which region this radial element belongs to
                    found_region = None
                    for region_id in self.solver.unique_regions:
                        # Check all connection types for this region
                        if (radial_idx in
                            self.solver.Index_radial_connections_region.get(
                                region_id, []
                            )):
                            found_region = region_id
                            break
                        elif (radial_idx in
                            self.solver.Index_normal_connections_region.get(
                                region_id, []
                            )):
                            found_region = region_id
                            break
                        elif (radial_idx in
                            self.solver.Index_struct_connections_region.get(
                                region_id, []
                            )):
                            found_region = region_id
                            break

                    if found_region is not None:
                        # Get the resistor material for this region
                        mat = self.solver.Conductor_resistor[found_region]

                        # Calculate resistivity at this T and B
                        rho = (
                            self.solver.materials.resistivity_selector(
                                mat,
                                np.array([temp]),
                                np.array([B_fixed]),
                                reshape=True
                            )
                        )
                        # Handle scalar or array return
                        if np.isscalar(rho):
                            rho = rho
                        else:
                            rho = rho[0]

                        # Get reference radial resistance (geometric part)
                        R_ref = self.solver.radial_resistance[radial_idx]
                        # Scale by temperature-dependent resistivity
                        # R = R_ref * (rho / rho_ref)
                        rho_ref = (
                            self.solver.ref_resistivity[radial_idx]
                            .cpu().numpy()
                        )
                        R_radial = R_ref * (rho / rho_ref)
                        R_total += R_radial

                R_circuit[idx] = R_total

            # Calculate time constant tau = L/R
            tau = L_circuit / R_circuit

            # Plot
            fig = plt.figure(figsize=(8, 6))
            ax = fig.add_subplot(111)
            ax.plot(T, tau, linewidth=2, color='blue')
            ax.set_xlabel('Temperature, K')
            ax.set_ylabel('Time Constant, s')
            ax.set_title(
                f'Circuit {circuit_id} Time Constant vs '
                f'Temperature (B = {B_fixed} T)'
            )
            ax.grid(True, alpha=0.3)
            ax.set_ylim(bottom=0)

            # Save figure
            fig.savefig(
                self.solver.sim_output + os.sep
                + f'TimeConstant_vs_T_Circuit_{circuit_id}.png',
                dpi=300
            )
            plt.close(fig)

    def generate_time_constant_vs_field(self):
        """
        Generates graphs of magnet time constant (tau = L/R) as a function
        of magnetic field at constant temperature for each circuit.

        For No-Insulation coils, the time constant is determined by the
        radial (turn-to-turn) resistance. Field dependence comes from
        magnetoresistance effects in conductor materials.
        :return:
        """
        # Magnetic field range for the sweep
        B = np.linspace(0.1, 50.0, 200)  # T (avoid B=0 for numerical stability)
        T_fixed = 5.0  # Fixed temperature in K

        # Create figure for each circuit
        for circuit_id in self.solver.unique_circuits:
            if circuit_id not in self.solver.inductance_dict:
                continue
            if circuit_id not in self.solver.index_dict['circuits']:
                continue

            # Get circuit inductance (constant)
            L_circuit = self.solver.inductance_dict[circuit_id]

            # Get radial connection indices for this circuit
            radial_indices = (
                self.solver.index_dict['circuits'][circuit_id]['IR']
            )

            # Skip if no radial elements
            if len(radial_indices) == 0:
                continue

            # Calculate total circuit resistance as function of field
            R_circuit = np.zeros_like(B)

            for idx, field in enumerate(B):
                R_total = 0.0

                # Sum radial resistance of all elements in this circuit
                for radial_idx in radial_indices:
                    # Find which region this radial element belongs to
                    found_region = None
                    for region_id in self.solver.unique_regions:
                        # Check all connection types for this region
                        if (radial_idx in
                            self.solver.Index_radial_connections_region.get(
                                region_id, []
                            )):
                            found_region = region_id
                            break
                        elif (radial_idx in
                            self.solver.Index_normal_connections_region.get(
                                region_id, []
                            )):
                            found_region = region_id
                            break
                        elif (radial_idx in
                            self.solver.Index_struct_connections_region.get(
                                region_id, []
                            )):
                            found_region = region_id
                            break

                    if found_region is not None:
                        # Get the resistor material for this region
                        mat = self.solver.Conductor_resistor[found_region]

                        # Calculate resistivity at this T and B
                        rho = (
                            self.solver.materials.resistivity_selector(
                                mat,
                                np.array([T_fixed]),
                                np.array([field]),
                                reshape=True
                            )
                        )
                        # Handle scalar or array return
                        if np.isscalar(rho):
                            rho = rho
                        else:
                            rho = rho[0]

                        # Get reference radial resistance (geometric part)
                        R_ref = self.solver.radial_resistance[radial_idx]
                        # Scale by temperature-dependent resistivity
                        # R = R_ref * (rho / rho_ref)
                        rho_ref = (
                            self.solver.ref_resistivity[radial_idx]
                            .cpu().numpy()
                        )
                        R_radial = R_ref * (rho / rho_ref)
                        R_total += R_radial

                R_circuit[idx] = R_total

            # Calculate time constant tau = L/R
            tau = L_circuit / R_circuit

            # Plot
            fig = plt.figure(figsize=(8, 6))
            ax = fig.add_subplot(111)
            ax.plot(B, tau, linewidth=2, color='blue')
            ax.set_xlabel('Magnetic Field, T')
            ax.set_ylabel('Time Constant, s')
            ax.set_title(
                f'Circuit {circuit_id} Time Constant vs '
                f'Magnetic Field (T = {T_fixed} K)'
            )
            ax.grid(True, alpha=0.3)
            ax.set_ylim(bottom=0)

            # Save figure
            fig.savefig(
                self.solver.sim_output + os.sep
                + f'TimeConstant_vs_B_Circuit_{circuit_id}.png',
                dpi=300
            )
            plt.close(fig)
    
    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)

        E_dis =     integrate.cumulative_trapezoid(np.sum(self.solver.P_AC_array, axis=1), self.solver.t_array, initial=0)
        E_dis +=    integrate.cumulative_trapezoid(np.sum(self.solver.Prad_array, axis=1), self.solver.t_array, initial=0)
        E_dis +=    integrate.cumulative_trapezoid(np.sum(self.solver.Pin_array, axis=1), self.solver.t_array, initial=0)

        ax1.plot(self.solver.t_array, self.solver.E_array)
        ax2.plot(self.solver.t_array, self.solver.Bz_center_array)
        ax3.plot(self.solver.t_array, np.sum(self.solver.P_AC_array, axis=1), label='AC loss')
        ax3.plot(self.solver.t_array, np.sum(self.solver.Prad_array, axis=1), label='Radial loss')
        ax3.plot(self.solver.t_array, np.sum(self.solver.Pin_array, axis=1), label='Input power')
        [ax4.plot(self.solver.t_array, np.array(self.solver.I_circuit_array)[:, i]) for i in range(len(self.solver.I_circuit_array[0]))]
        [ax5.plot(self.solver.t_array, np.array(self.solver.V_circuit_array)[:, i]) for i in range(len(self.solver.V_circuit_array[0]))]
        ax6.plot(self.solver.t_array, self.solver.solver_t_array)
        ax7.plot(self.solver.t_array, 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('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')
        ax3.legend()

        fig.tight_layout()

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

    def plot_ramp_rate_interpolation(self):
        for i in self.solver.unique_circuits:
            fig = plt.figure(figsize=(10, 5))
            ax1 = fig.add_subplot(121)
            ax2 = fig.add_subplot(122)
            t = np.linspace(self.solver.t0[0], self.solver.t0[1], 500000)
            dIdt = self.solver.circuit_interp_dIdt[i](t)
            I = integrate.cumulative_trapezoid(dIdt, t, initial=0)

            ax1.plot(t, dIdt, label='Circuit: ' + str(i))
            ax1.set_xlabel('Time, s')
            ax1.set_ylabel('dIdt, A/s')

            ax2.plot(t, I, label='Circuit: ' + str(i))
            ax2.set_xlabel('Time, s')
            ax2.set_ylabel('Current, A')

            fig.savefig(self.solver.sim_output + os.sep + 'dIdt_' + str(i) + '.png', dpi=300)
            plt.close(fig)

    def plot_time_step_interpolation(self, t_densified, t_diff):
        fig = plt.figure(figsize=(10, 5))
        ax1= fig.add_subplot(111)

        t = np.linspace(self.solver.t0[0], self.solver.t0[1], 10000)
        dt = np.zeros_like(t)
        for i in range(len(t)):
            dtt = self.solver.time_step_dt(t[i]).item() # Interpolated timestep
            dtu = np.min(self.solver.unique_sorted_timesteps[self.solver.unique_sorted_timesteps > t[i]]) - t[i] if t[i] < np.max(self.solver.unique_sorted_timesteps) else 1e99 # Upper bound time step.
            dt[i] = np.min([dtt, dtu, self.solver.max_timestep])

        ax1.plot(t, dt)
        ax1.plot(t_densified, t_diff)
        ax1.set_xlabel('Time, s')
        ax1.set_ylabel('dt A/s')
        fig.savefig(self.solver.sim_output + os.sep + 'dt.png', dpi=300)
        plt.close()


class SummaryPlotter:
    """
    A class for creating summary plots from simulation results.
    
    This plotter creates dual-subplot figures showing time-series data from 
    simulation summary CSV files. It supports configurable parameters,
    labels, and plot styling options.
    """
    
    def __init__(self, config_path: Optional[str] = None, **kwargs):
        """
        Initialize the SummaryPlotter.
        
        Args:
            config_path: Path to YAML configuration file
            **kwargs: Direct configuration parameters (override YAML config)
        """
        # Default configuration
        self.config = {
            'input_paths': [],
            'labels': [],
            'parameters': ['T', 'Fr_density'],
            'output_filename': 'summary_plot.png',
            'figure_size': (16, 6),
            'layout': 'horizontal',  # 'horizontal' or 'vertical'
            'dpi': 1000,
            'time_limits': {'start': 0.1, 'end': 0.5},
            'force_limits': {'min': 0.0, 'max': 70},
            'temperature_limits': {'min': 0.0, 'max': 300},
            'tag': None,
            'colormap': 'coolwarm',
            'show_plot': True
        }
        
        # Load configuration from YAML if provided
        if config_path and os.path.exists(config_path):
            self._load_config(config_path)
        
        # Override with any direct parameters
        self.config.update(kwargs)
        
        # Initialize plot elements
        self.fig = None
        self.ax1 = None
        self.ax2 = None
        
        # Set up matplotlib parameters
        self._setup_matplotlib_params()
    
    def _load_config(self, config_path: str) -> None:
        """Load configuration from YAML file."""
        try:
            with open(config_path, 'r') as f:
                yaml_config = yaml.safe_load(f)
                self.config.update(yaml_config)
        except Exception as e:
            print(f"Warning: Could not load config file {config_path}: {e}")
    
    def _setup_matplotlib_params(self) -> None:
        """Configure matplotlib parameters for consistent styling."""
        params = {
            'axes.labelsize': 11,
            'xtick.labelsize': 11,
            'ytick.labelsize': 11,
            'legend.fontsize': 11,
            'font.sans-serif': 'Arial',
            'text.usetex': False,
            'figure.autolayout': True,
            'lines.markersize': 5,
            'axes.grid': True,
            'grid.color': 'lightGrey',
            'axes.axisbelow': True
        }
        
        for key, value in params.items():
            rcParams[key] = value
    
    def _create_figure(self) -> None:
        """Create figure and subplots based on configuration."""
        figsize = tuple(self.config['figure_size'])
        self.fig = plt.figure(figsize=figsize)
        
        if self.config['layout'] == 'horizontal':
            self.ax1 = self.fig.add_subplot(121)
            self.ax2 = self.fig.add_subplot(122)
            self.fig.subplots_adjust(wspace=0.035)
        else:  # vertical layout
            self.ax1 = self.fig.add_subplot(211)
            self.ax2 = self.fig.add_subplot(212)
    
    def _setup_axis_labels(self, parameter: str) -> tuple:
        """
        Set up axis labels based on the parameter type.
        
        Args:
            parameter: The parameter being plotted
            
        Returns:
            tuple: (ylabel_ax1, ylabel_ax2, column_name)
        """
        label_map = {
            'Fr': ('Peak Force, kN', None, 'Max Fr'),
            'Fz': ('Peak Force, kN', None, 'Max Fz'),
            'Fr_density': ('Peak Force Density, N/mm$^3$', None, 'Max Fr Density'),
            'Fz_density': ('Peak Force Density, N/mm$^3$', None, 'Max Fz Density'),
            'I_density': ('Peak Current Density, A/mm$^2$', None, 'Max I Density'),
            'II': ('Current, kA', None, 'Max II'),
            'IR': ('Current, kA', None, 'Max IR'),
            'T': (None, 'Temperature, K', 'maxT'),
            'Bz': ('Magnetic Field, T', None, 'Max Bz'),
            'Br': ('Magnetic Field, T', None, 'Max Br')
        }
        
        return label_map.get(parameter, (None, None, parameter))
    
    def _plot_data(self, data_list: List[pd.DataFrame]) -> None:
        """
        Plot the data on the configured axes.
        
        Args:
            data_list: List of pandas DataFrames containing the data
        """
        # Set up colormap
        cmap = matplotlib.colormaps[self.config['colormap']]
        colors = [cmap(i) for i in np.linspace(0, 1.0, len(data_list))]
        
        # Plot each parameter
        for param in self.config['parameters']:
            ylabel_ax1, ylabel_ax2, column_name = self._setup_axis_labels(param)
            
            # Set axis labels
            if ylabel_ax1 and self.ax1:
                self.ax1.set_ylabel(ylabel_ax1)
            if ylabel_ax2 and self.ax2:
                self.ax2.set_ylabel(ylabel_ax2)
            
            # Plot data for each input path
            for i, data in enumerate(data_list):
                if column_name not in data.columns:
                    print(f"Warning: Column '{column_name}' not found in data {i}")
                    continue
                
                # Get time reference (start time)
                t0 = data['time'].values[0]
                x = data['time'].values
                y = data[column_name].values
                
                # Apply scaling and plot on appropriate axis
                label = self.config['labels'][i] if i < len(self.config['labels']) else f'Data {i}'
                
                if param in ['Fr_density', 'Fz_density']:
                    # Force density: scale and plot on ax1
                    if self.ax1:
                        self.ax1.plot(x - t0, y * 1e-9, color=colors[i], label=label)
                else:
                    # Temperature and other parameters: plot on ax2
                    if self.ax2:
                        self.ax2.plot(x - t0, y, color=colors[i], label=label)
    
    def _apply_styling(self) -> None:
        """Apply styling, limits, and formatting to the plots."""
        # Set x-axis properties
        time_limits = self.config['time_limits']
        for ax in [self.ax1, self.ax2]:
            if ax:
                ax.set_xlabel('Time, s')
                ax.set_xlim((time_limits['start'], time_limits['end']))
        
        # Set y-axis limits if configured
        if 'force_limits' in self.config and self.ax1:
            force_limits = self.config['force_limits']
            self.ax1.set_ylim((force_limits['min'], force_limits['max']))

        # Set y-axis limits if configured
        if 'temperature_limits' in self.config and self.ax1:
            temperature_limits = self.config['temperature_limits']
            self.ax2.set_ylim((temperature_limits['min'], temperature_limits['max']))
        
        # Add legends
        if self.ax1:
            self.ax1.legend(loc='best', ncols=2)
        if self.ax2:
            self.ax2.legend(loc='lower right', ncols=2)
        
        # Apply tick formatting
        self._format_ticks()
    
    def _format_ticks(self) -> None:
        """Apply consistent tick formatting to both axes."""
        tick_params = {
            'direction': 'in',
            'width': 1.2,
            'length': 5
        }
        
        minor_tick_params = {
            'direction': 'in',
            'width': 1.0,
            'length': 3,
            'which': 'minor'
        }
        
        # Format ax1 ticks
        if self.ax1:
            self.ax1.xaxis.set_tick_params(top=True, **tick_params)
            self.ax1.yaxis.set_tick_params(left=True, right=True, **tick_params)
            self.ax1.xaxis.set_tick_params(top=True, **minor_tick_params)
            self.ax1.yaxis.set_tick_params(left=True, right=True, **minor_tick_params)
            
            # Set tick locators
            self.ax1.xaxis.set_minor_locator(MultipleLocator(0.05))
            self.ax1.xaxis.set_major_locator(MultipleLocator(0.1))
            self.ax1.yaxis.set_minor_locator(MultipleLocator(5.0))
            self.ax1.yaxis.set_major_locator(MultipleLocator(10.0))
        
        # Format ax2 ticks
        if self.ax2:
            self.ax2.xaxis.set_tick_params(top=True, **tick_params)
            self.ax2.yaxis.set_tick_params(left=True, right=True, **tick_params)
            self.ax2.xaxis.set_tick_params(top=True, **minor_tick_params)
            self.ax2.yaxis.set_tick_params(left=True, right=True, **minor_tick_params)
            
            # Set tick locators
            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))
    
    def create_plot(self) -> None:
        """
        Create the summary plot based on the current configuration.
        
        This method loads data from the specified input paths, creates the figure,
        plots the data, applies styling, and saves/shows the result.
        """
        # Validate configuration
        if not self.config['input_paths']:
            raise ValueError("No input paths specified")
        
        # Load data from all input paths
        data_list = []
        for path in self.config['input_paths']:
            summary_path = os.path.join(path, 'summary.csv')
            if not os.path.exists(summary_path):
                print(f"Warning: Summary file not found at {summary_path}")
                continue
            
            try:
                data = pd.read_csv(summary_path)
                data_list.append(data)
            except Exception as e:
                print(f"Error loading data from {summary_path}: {e}")
        
        if not data_list:
            raise ValueError("No valid data files found")
        
        # Create figure and plot data
        self._create_figure()
        self._plot_data(data_list)
        self._apply_styling()
        
        # Finalize plot
        self.fig.tight_layout()
        
        # Save plot
        output_path = self.config['output_filename']
        self.fig.savefig(output_path, dpi=self.config['dpi'])
        print(f"Plot saved to: {output_path}")
        
        # Show plot if requested
        if self.config['show_plot']:
            plt.show()
        else:
            plt.close(self.fig)
    
    @classmethod
    def from_yaml(cls, config_path: str) -> 'SummaryPlotter':
        """
        Create a SummaryPlotter instance from a YAML configuration file.
        
        Args:
            config_path: Path to the YAML configuration file
            
        Returns:
            SummaryPlotter instance
        """
        return cls(config_path=config_path)


# Legacy alias for backward compatibility
plotter_v3 = SummaryPlotter
