from refDLL import RefProp, RegDllCall
import numpy as np
import matplotlib.pyplot as plt
from plotly import graph_objs as go
import pandas as pd
import altair as alt
alt.data_transformers.disable_max_rows()

# Setting default font sizes for plots
SMALL_SIZE = 10
MEDIUM_SIZE = 22
BIGGER_SIZE = 28

class Diagram_PH:
    """Class to define and plot PH diagrams for a specified refrigerant."""
    
    def __init__(self, REFRIG):
        """
        Initialize the Diagram_PH class with a specific refrigerant.
        
        Args:
            REFRIG (str): The name of the refrigerant.
        """
        self.Refname = REFRIG  # Name of the refrigerant
        self.callref = RegDllCall(self.Refname)  # Register DLL call with the refrigerant name
        self.Hsl, self.Hsv, self.Psat, self.Tsat = self.get_psat_values()  # Get saturation values
        self.Tmax, self.Tmin, self.T_lst, self.P, self.IsoT_lst = self.get_IsoT_values()  # Get isothermal values
        self.extra_points = []  # List for additional points to be plotted
        self.extra_points_order = []  # List to store the order of the extra points
        self.extra_dict = {}  # Dictionary to store extra points by order
        self.nodes = []  # List for node annotations in Plotly plots

    def clearAllExtraPoint(self):
        """Clear all extra points previously added to the diagram."""
        self.extra_points = [] 
        self.extra_points_order = []
        self.extra_dict = {}
        
    def get_psat_values(self):
        """
        Calculate the psat values for the refrigerant.

        Returns:
            tuple: The Hsl, Hsv, Psat, and Tsat values.
        """   
        Hsl, Hsv, Psat, Tsat = [], [], [], []
        
        # Calculate values for different pressures in the range of the refrigerant's pressure
        for p in np.arange(self.callref.refrig.p_begin(), self.callref.refrig.p_end(), 0.5e5):
            # Calculate and append the liquid enthalpy for the given pressure
            Hsl.append(self.callref.refrig.hsl_px(p, 0) / 1e3)
            # Calculate and append the vapor enthalpy for the given pressure
            Hsv.append(self.callref.refrig.hsv_px(p, 1) / 1e3)
            # Append the pressure
            Psat.append(p / 1e5)
            # Calculate and append the saturation temperature for the given pressure
            Tsat.append(self.callref.refrig.T_px(p, 0.5))
            
            # Stop calculation if the liquid enthalpy doesn't change anymore
            if len(Hsl) > 2 and Hsl[-1] == Hsl[-2]:
                break

        return Hsl, Hsv, Psat, Tsat
    
    def add_points_common(self, refppt, points, is_ordered=False):
        """
        Add extra points to the diagram.

        Args:
            refppt (int): The property pair identifier.
            points (dict): The points to be added.
            is_ordered (bool, optional): Whether the points are ordered. Defaults to False.
        """
                
        # Mapping for the h calculation functions
        h_calc_funcs = {
            RefProp.PX: lambda p, x: self.callref.refrig.h_px(p, x),
            RefProp.PT: self.callref.H_pT,
            RefProp.TSX: lambda t, x: self.callref.h_px(self.callref.p_Tx(t, round(x)), x),
            RefProp.TSSH: lambda t, x: self.callref.h_px(self.callref.p_Tx(t, round(x)), x),
        }
        
        # Mapping for the p calculation functions
        p_calc_funcs = {
            RefProp.PX: lambda p, _: p,
            RefProp.PT: lambda p, _: p,
            RefProp.TSX: lambda t, x: self.callref.p_Tx(t, round(x)),
            RefProp.TSSH: lambda t, x: self.callref.p_Tx(t, round(x)),
        }
        
        # Iterate over points
        extra_dict = {}
        for _, i in enumerate(points):
            point = points[i]
            # Calculate h and p values using the corresponding function
            h = h_calc_funcs[refppt](*point)
            p = p_calc_funcs[refppt](*point)
            if is_ordered:
                extra_dict[i] = (h * 1e-3, p * 1e-5)  # Use index as order
            else:
                # If the points are not ordered, simply append them to the list
                self.extra_points.append([h * 1e-3, p * 1e-5])
        # If the points are ordered, store them in the dictionary using the index as the order        
        if is_ordered:
            self.extra_dict.update(extra_dict)
        
    def add_points(self, data):
        """
        Add extra points to the diagram.

        Args:
            data (dict): The points to be added.
        """        
        for refppt, points in data.items():
            self.add_points_common(refppt, points, False)

    
    def add_points_order(self, data):
        """
        Add extra ordered points to the diagram.

        Args:
            data (dict): The points to be added.
        """          
        for refppt, points in data.items():
            self.add_points_common(refppt, points, True)
            self.extra_points_order = sorted(self.extra_dict.items(), key=lambda item: item[0])

    def get_IsoT_values(self):
        """
        Calculate the isothermal values for the refrigerant.

        Returns:
            tuple: The Tmax, Tmin, T_lst, P, and IsoT_lst values.
        """
                    
        # Calculate the temperatures for different pressures in the range of the refrigerant's pressure
        T = [self.callref.refrig.T_px(p, 0.5) - 273.15 for p in np.arange(self.callref.refrig.p_begin(), self.callref.refrig.p_end(), 50e5)]
        
        # Find the maximum and minimum saturation temperatures
        Tmax, Tmin = max(self.Tsat) - 273.15 - 1, min(self.Tsat) - 273.15
        
        # Find the list of temperatures to use for isothermal calculations
        T_lst = self.callref.findwhole10number(Tmin, Tmax)
        
        # Generate pressures to use for isothermal calculations
        P = np.arange(self.callref.refrig.p_begin(), self.callref.refrig.p_end(), 0.05e5)
        
        # Calculate isothermal values for each temperature in the list
        IsoT_lst = [[self.callref.refrig.h_pT(p, temp + 273.15) / 1e3 for p in P] for temp in T_lst]

        data = {
            'Temperature': T_lst.repeat(len(P)),
            'Pressure': np.tile(P, len(T_lst)),
            'Enthalpy': np.concatenate(IsoT_lst)
        }
    
        df = pd.DataFrame(data)
    
        # Save the dataframe to a CSV file for later analysis
        # df.to_csv(r'C:\Users\serameza\impact\EMEA_MBD_GitHub\CheckLabdata\IsothermalData.csv', index=False)
        return Tmax, Tmin, T_lst, P, IsoT_lst

    def add_points_df(self, df):
        """
        Add points to the diagram from a DataFrame, considering groups and including the node index.
        
        Args:
            df (DataFrame): DataFrame containing the data points to be added.
        """
        df_sorted = df.sort_values(by='Order')
        for idx, row in df_sorted.iterrows():
            # Include 'Node' from the DataFrame index in the extra_dict
            self.extra_dict[row['Order']] = {
                'Enthalpy': row['Enthalpy'] / 1e3,  # Convert to kJ/kg if originally in J/kg
                'Pressure': row['Pressure'] / 1e5,  # Convert to bar if originally in Pa
                'Group': row.get('Group'),
                'Node': idx  # Assuming 'idx' is the node index from the original DataFrame
            }
        # Sort the extra points by their order for later plotting
        self.extra_points_order = sorted(self.extra_dict.items(), key=lambda item: item[0])

    def plot_diagram(self):
        """Plot the PH diagram using Matplotlib."""
        plt.rc('font', size=SMALL_SIZE)          # Controls default text sizes
        plt.rc('axes', titlesize=SMALL_SIZE)     # Font size of the axes title
        plt.rc('axes', labelsize=MEDIUM_SIZE)    # Font size of the x and y labels
        plt.rc('xtick', labelsize=SMALL_SIZE)    # Font size of the tick labels
        plt.rc('ytick', labelsize=SMALL_SIZE)    # Font size of the tick labels
        plt.rc('legend', fontsize=SMALL_SIZE)    # Legend font size
        plt.rc('figure', titlesize=BIGGER_SIZE)  # Font size of the figure title
        plt.figure(figsize=[15, 10])

        # Plot saturation lines
        plt.plot(self.Hsl, self.Psat, 'k-', label='Liquid Saturation')
        plt.plot(self.Hsv, self.Psat, 'k-', label='Vapor Saturation')

        # Plot isotherms
        for Th_lst, temp in zip(self.IsoT_lst, self.T_lst):
            plt.plot(Th_lst, self.P / 1e5, 'g--', label=f'{temp}°C Isotherm', alpha=0.5)
            plt.annotate('{:.0f}°C'.format(temp), 
                         (self.callref.refrig.h_px(self.callref.refrig.p_Tx(temp + 273.15, 0.5), 0.1) / 1e3, self.callref.refrig.p_Tx(temp + 273.15, 0.5) / 1e5),
                         ha='center', 
                         backgroundcolor="white")  
            
        plt.yscale('log')
        
        # Plot additional points, grouped and connected by lines if applicable
        if self.extra_points_order:
            # Extract the groups and points for plotting
            df = pd.DataFrame(self.extra_points_order, columns=['Order', 'Data'])
            df['Enthalpy'] = df['Data'].apply(lambda x: x['Enthalpy'])
            df['Pressure'] = df['Data'].apply(lambda x: x['Pressure'])
            df['Group'] = df['Data'].apply(lambda x: x.get('Group', 'Unspecified'))
            plt.plot(df['Enthalpy'], df['Pressure'], '-o', zorder=10)
            # Plot points by group and connect them
            for group, group_df in df.groupby('Group'):
                if group != 'Unspecified':  # Only plot specified groups with lines
                    group_df = group_df.sort_values(by='Pressure')
                    plt.plot(group_df['Enthalpy'], group_df['Pressure'], '-o', zorder=10)

        plt.xlabel('Enthalpy [kJ/kg]')
        plt.ylabel('Pressure [bar]')
        plt.title(f'PH Diagram for {self.Refname}')
        plt.grid(True, which='both', linestyle='--')
        plt.tight_layout()
        return plt

    def plot_diagram_plotly(self):
        """Plot the PH diagram interactively using Plotly, with points connected by group."""
        fig = go.Figure()

        # Saturation lines
        fig.add_trace(go.Scatter(x=self.Hsl, y=self.Psat, mode='lines', name='Liquid Saturation', line=dict(color='black')))
        fig.add_trace(go.Scatter(x=self.Hsv, y=self.Psat, mode='lines', name='Vapor Saturation', line=dict(color='black')))

        # Isotherms
        for Th_lst, temp in zip(self.IsoT_lst, self.T_lst):
            fig.add_trace(go.Scatter(x=Th_lst, y=self.P / 1e5, mode='lines', name=f'{temp}°C Isotherm', line=dict(color='green', dash='dash', width=0.5)))

            # Add annotation for each isotherm
            enthalpy_at_mid_pressure = self.callref.refrig.h_px(self.callref.refrig.p_Tx(temp + 273.15, 0.5), 0.1) / 1e3
            pressure_at_mid_point = self.callref.refrig.p_Tx(temp + 273.15, 0.5) / 1e5
            pressure_at_mid_point = np.log10(pressure_at_mid_point)
            fig.add_annotation(
                x=enthalpy_at_mid_pressure, 
                y=pressure_at_mid_point,
                text=f'{temp:.0f}°C',
                showarrow=False,
                bgcolor="white"
            )

        if self.extra_points_order:
            # Prepare a DataFrame for easier handling
            df = pd.DataFrame([(order, data) for order, data in self.extra_points_order], columns=['Order', 'Data'])
            df['Enthalpy'] = df['Data'].apply(lambda x: x['Enthalpy'])
            df['Pressure'] = df['Data'].apply(lambda x: x['Pressure'])
            df['Group'] = df['Data'].apply(lambda x: x.get('Group', 'Unspecified'))
            df['Node'] = df['Data'].apply(lambda x: x['Node'])  # Assuming 'self.nodes' are in the same order as 'self.extra_points_order'
            
            fig.add_trace(go.Scatter(x=df['Enthalpy'], y=df['Pressure'], mode='markers+lines', name='Ordered Points',
                                     line=dict(color='red'),
                                     hoverinfo='text',
                                     text=df['Node']))  
        # Plot points by group
            for _, group_df in df.groupby('Group'):
                fig.add_trace(go.Scatter(
                    x=group_df['Enthalpy'],
                    y=group_df['Pressure'],
                    line=dict(color='red'),
                    hoverinfo='text',
                    text=group_df['Node'],
                    mode='markers+lines'
                ))

        # Update layout for readability
        fig.update_layout(
            xaxis=dict(
                title='Enthalpie [kJ/kg]',  # Title of x-axis
                showgrid=True,  # Show grid
                gridcolor='LightPink',  # Grid color
                linecolor='black',  # Axis line color
                linewidth=0.3,  # Axis line width
                mirror=True,  # Mirror axis lines
            ),
            yaxis=dict(
                title='Pression [bar]',  # Title of y-axis
                type='log',  # Use logarithmic scale
                showgrid=True,  # Show grid
                gridcolor='LightBlue',  # Grid color
                linecolor='black',  # Axis line color
                linewidth=0.3,  # Axis line width
                mirror=True,  # Mirror axis lines
            ),
            showlegend=False,  # Hide legend
            autosize=True,
            width=1000,
            height=800,
            margin=dict(l=100, r=50, b=100, t=100, pad=4),
            plot_bgcolor="white",
        )

        return fig

    def plot_diagram_altair(self):
        """Plot the PH diagram using Altair."""
        # Convert lists to DataFrame for Altair
        data_saturationL = pd.DataFrame({
            'Enthalpy': self.Hsl,
            'Pressure': self.Psat,
            'Type': ['Liquid Saturation'] * len(self.Hsl)
        })

        data_saturationV = pd.DataFrame({
            'Enthalpy': self.Hsv,
            'Pressure': self.Psat,
            'Type': ['Liquid Saturation'] * len(self.Hsv)
        })

        # Isotherms and annotations
        data_isotherms = pd.DataFrame({
            'Enthalpy': np.concatenate(self.IsoT_lst),
            'Pressure': np.tile(self.P / 1e5, len(self.T_lst)),
            'Temperature': np.repeat(self.T_lst, len(self.P))
        })
        df_extra = pd.DataFrame()
        # Additional points, if present
        if self.extra_points_order:
            # Prepare a DataFrame for easier handling
            df = pd.DataFrame([(order, data) for order, data in self.extra_points_order], columns=['Order', 'Data'])
            
            df_extra['Enthalpy'] = df['Data'].apply(lambda x: x['Enthalpy'])
            df_extra['Pressure'] = df['Data'].apply(lambda x: x['Pressure'])
            df_extra['Group'] = df['Data'].apply(lambda x: x.get('Group', 'Unspecified'))
            df_extra['Node'] = df['Data'].apply(lambda x: x['Node'])  # Assuming 'self.nodes' are in the same order as 'self.extra_points_order'
            df_extra['Order'] = df.index.to_list()
        else:
            df_extra = pd.DataFrame(columns=['Enthalpy', 'Pressure', 'Node'])

        # Create the base chart
        base = alt.Chart().encode(
            x=alt.X('Enthalpy:Q', title='Enthalpie [kJ/kg]'),
            y=alt.Y('Pressure:Q', title='Pression [bar]', scale=alt.Scale(type='log'))
        )

        # Liquid saturation chart
        chart_saturationL = alt.Chart(data_saturationL).mark_line().encode(
            x='Enthalpy:Q',
            y=alt.Y('Pressure:Q', scale=alt.Scale(type='log')),
        )

        # Vapor saturation chart
        median_pressure = data_saturationV['Pressure'].median()

        # Split DataFrame into two parts
        data_upper = data_saturationV[data_saturationV['Pressure'] > median_pressure]
        data_lower = data_saturationV[data_saturationV['Pressure'] <= median_pressure]

        # Create and combine charts
        chart_upper = alt.Chart(data_upper).mark_point(filled=True, size=2).encode(
            x='Enthalpy:Q',
            y=alt.Y('Pressure:Q', scale=alt.Scale(type='log'))
        )

        chart_lower = alt.Chart(data_lower).mark_line().encode(
            x='Enthalpy:Q',
            y=alt.Y('Pressure:Q', scale=alt.Scale(type='log'))
        )

        chart_saturationV = chart_upper + chart_lower 
  
        data_isotherms.sort_values(by=['Enthalpy'], inplace=True, ascending=[False])
        data_isotherms['Order'] = data_isotherms.index.to_list()

        # Isotherms chart
        chart_isotherms = alt.Chart(data_isotherms).mark_line(opacity=0.5).encode(
            x='Enthalpy:Q',
            y='Pressure:Q',
            order='Order:N',
            color=alt.Color('Temperature:Q', legend=alt.Legend(title="Temperature (°C)"))
        )
        
        # Add annotations for isotherms (Altair does not handle annotations directly like Plotly)
        df_extra['Group'] = df_extra['Group'].fillna(-2000).astype(str)
        # Additional points
        grouped = df_extra[df_extra['Group'] != -2000]
        grouped['Type'] = grouped['Order']
        brush = alt.selection_interval()

        # Safety check before using the DataFrame
        if isinstance(grouped, pd.DataFrame) and 'Group' in grouped.columns:
            group_chart = alt.Chart(grouped).mark_line(point=True).encode(
                x='Enthalpy:Q',
                y='Pressure:Q',
                detail='Group:N',
                color=alt.value('red'),
                tooltip=['Node:N', 'Enthalpy:Q', 'Pressure:Q']
            ).add_params(
                brush
            )
        else:
            print("Error: 'grouped' is not a DataFrame or missing necessary columns")

        # Similar check for 'df_extra'
        if isinstance(df_extra, pd.DataFrame) and {'Enthalpy', 'Pressure', 'Node', 'Order'}.issubset(df_extra.columns):
            ungroup_chart = alt.Chart(df_extra).mark_line(point=True).encode(
                x='Enthalpy:Q',
                y=alt.Y('Pressure:Q', scale=alt.Scale(type='log')),
                tooltip=['Node:N', 'Enthalpy:Q', 'Pressure:Q'],
                order='Order:O'
            ).add_params(
                brush
            )
        else:
            print("Error: 'df_extra' is not a DataFrame or missing necessary columns")

        # Combine charts only if both are defined
        if group_chart and ungroup_chart:
            chart_extra_points = group_chart + ungroup_chart
        else:
            chart_extra_points = group_chart if group_chart else ungroup_chart

        # Combine all charts
        final_chart = chart_saturationL + chart_saturationV + chart_isotherms
        if chart_extra_points is not None:
            final_chart += chart_extra_points

        # Final configuration and display
        final_chart = final_chart.properties(
            width=800,
            height=600
        ).configure_axis(
            grid=True
        ).configure_view(
            strokeWidth=0
        )
        interactive_scatter = final_chart.encode().brush(
            alt.selection_interval()
        ).interactive()
        return interactive_scatter