import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
# try:
#     import Thermobar as pt
# except:
#     print('Thermobar cannot be imported, please check your numpy version')

from petthermotools.GenFuncs import *
from petthermotools.GenFuncs import _ensure_julia_ready
from petthermotools.Plotting import *
from petthermotools.Liq import *
from petthermotools.MELTS import *
from petthermotools.Path import *
import multiprocessing
from multiprocessing import Queue
from multiprocessing import Process
import time
from scipy import interpolate
# from shapely.geometry import MultiPoint, Point, Polygon

import numpy as np
import pandas as pd
import multiprocessing
import time
from pathlib import Path

def path_4_saturation_multi(q, index, *, Model = None, P_bar = None, comp = None, T_maxdrop_C = None, dt_C = None, T_initial_C = None, fO2_buffer = None,
                      fO2_offset = 0.0, H2O_Sat = None, phases = None):
    """
    Worker function to run crystallization simulations at a subset of pressures in parallel
    using alphaMELTS for Python or MAGEMinCalc to determine mineral co-saturations.

    Parameters
    ----------
    q : multiprocessing.Queue
        Queue object for inter-process communication (used to return results).
    index : list of int or numpy.ndarray
        Chunk of indices from the full pressure array (`P_bar`) assigned to this worker process.
    Model : str
        Thermodynamic model to use: e.g., 'MELTSv1.0.2', 'pMELTS', or 'MAGEMinCalc'.
    P_bar : numpy.ndarray
        The full array of pressures (in bar) that is being iterated over by all workers.
    comp : dict
        Bulk composition (oxides in wt%) of the starting material.
    T_maxdrop_C : float
        The minimum temperature (°C) below the liquidus to continue the simulation path down to.
    dt_C : float
        Step size for temperature change (in °C).
    T_initial_C : float
        Starting temperature (°C) for liquidus calculations.
    fO2_buffer : str, optional
        Oxygen fugacity buffer ('FMQ' or 'NNO' for MELTS, 'qfm' or 'nno' for MAGEMinCalc).
    fO2_offset : float, default 0.0
        Offset in log units from the specified fO2 buffer.
    H2O_Sat : bool
        If True, the system is simulated under water-saturated conditions.
    phases : list of str
        List of phases (e.g., ['quartz1', 'alkali-feldspar1']) to monitor for saturation.

    Returns
    -------
    None
        The results are placed into the multiprocessing queue `q` as a two-element list:
        `[list of indices processed (int), dict of raw simulation results]`.
    """

    results = {}
    idx = []
    
    if "MELTS" in Model:
        from meltsdynamic import MELTSdynamic

        if Model is None or Model == "MELTSv1.0.2":
            melts = MELTSdynamic(1)
        elif Model == "pMELTS":
            melts = MELTSdynamic(2)
        elif Model == "MELTSv1.1.0":
            melts = MELTSdynamic(3)
        elif Model == "MELTSv1.2.0":
            melts = MELTSdynamic(4)
    else:
        from juliacall import Main as jl, convert as jlconvert
        env_dir = Path.home() / ".petthermotools_julia_env"
        jl_env_path = env_dir.as_posix()

        jl.seval(f"""
            import Pkg
            Pkg.activate("{jl_env_path}")
            """)

        jl.seval("using MAGEMinCalc")

        comp_julia = jl.seval("Dict")(comp)   

    for i in index:
        # try:
        if "MELTS" in Model:
            q.put([i, -1, "Start"])
            Results, tr = path_MELTS(P_bar = P_bar[i],
                    Model=Model,
                    comp=comp,
                    T_maxdrop_C=T_maxdrop_C,
                    dt_C=dt_C,
                    T_initial_C=T_initial_C,
                    find_liquidus=True,
                    fO2_buffer=fO2_buffer,
                    fO2_offset=fO2_offset,
                    fluid_sat=H2O_Sat,
                    Suppress=['rutile', 'tridymite'],
                    phases=phases,
                    trail = True,
                    melts = melts,
                    q = q,
                    index = i
                )
        else:
            if fO2_offset is None:
                fO2_offset = 0.0

            Results_df = jl.MAGEMinCalc.path(
                                    comp=comp_julia, dt_C=dt_C,
                                    P_bar=P_bar[i], T_min_C = T_maxdrop_C,            
                                    Model=Model, fo2_buffer=fO2_buffer, 
                                    fo2_offset=fO2_offset, find_liquidus=True,
                                    frac_xtal = False
                                    )
            
            Results = dict(Results_df)
            
            tr = True

        idx.append(i)


        results[f"Run {i}"] = Results

        if tr is False:
            break
        # except:
        #     idx.append(i)
    if "MELTS" not in Model:
        q.put([idx, results])
    return

def mineral_cosaturation(Model="MELTSv1.0.2", cores=int(np.floor(multiprocessing.cpu_count())), bulk=None,
                         phases=['quartz1', 'alkali-feldspar1', 'plagioclase1'],
                         P_bar=np.linspace(250, 5000, 32), Fe3Fet_init=None, H2O_init=None,
                         CO2_init=None, H2O_Sat=False, T_initial_C=None, dt_C=2,
                         T_maxdrop_C=50, T_cut_C=20, find_range=True,
                         find_min=True, fO2_buffer=None, fO2_offset=0.0, timeout = 90, multi_processing = True):
    """
    Determines the pressure at which two or more minerals co-saturate (following the method of Gualda et al. 2014).

    The analysis involves running multiple isobaric crystallization paths over a range of
    pressures, calculating the saturation temperature for the specified phases at each pressure,
    and finding the pressure that minimizes the temperature difference between them.

    Parameters
    ----------
    Model : str
        Thermodynamic model to use.
    cores : int
        Number of parallel processes to use.
    bulk : dict
        Bulk composition of the starting material, reported as oxide concentrations in wt%.
    phases : list
        Mineral phases to track for saturation. Must be a 2 or 3 item list ['quartz1', 'plagioclase1']
    P_bar : array-like
        Array of pressures used for the crystallization calculations.
    Fe3Fet_init : float, optional
        Initial Fe³⁺/Fe_total ratio.
    H2O_init : float, optional
        Initial H₂O content in the system (wt%).
    CO2_init : float, optional
        Initial CO₂ content in the system (wt%).
    H2O_Sat : bool
        If True, run at H₂O saturation. System must have enough H2O to be saturated.
    T_initial_C : float, optional
        Starting temperature in °C for liquidus calculations.
    dt_C : float
        Temperature step size (°C).
    T_maxdrop_C : float
        Maximum temperature drop below the liquidus to search for satuartion (°C).
    T_cut_C : float
        Maximum acceptable difference in phase saturation temperatures for "co-saturation" (°C). Used when trying to fit a polynomial to the results to find the true minimum point.
    find_range : bool
        If True, analyze phase co-saturation ranges. Currently disabled.
    find_min : bool
        If True, determine pressure where phase saturation temperature difference is minimized.
    fO2_buffer : str, optional
        Oxygen fugacity buffer (e.g., "FMQ", "NNO").
    fO2_offset : float, default 0.0
        Offset from the fO2 buffer in log units.
    timeout : int
        Timeout (in seconds) for each process.
    multi_processing : bool
        If True, run simulations in parallel using multiprocessing.

    Returns
    -------
    out : dict
        Contains either 'Output' (DataFrame of phase saturation conditions) or both 'Output' and 'CurveMin' (results from the minimization calculations).
    Results : dict
        Raw simulation outputs for each pressure step.
    """
    ## make sure everything is a float
    T_initial_C = to_float(T_initial_C)
    T_maxdrop_C = to_float(T_maxdrop_C)
    T_cut_C = to_float(T_cut_C)

    P_bar      = to_float(P_bar)

    Fe3Fet_init= to_float(Fe3Fet_init)
    H2O_init   = to_float(H2O_init)
    CO2_init   = to_float(CO2_init)
    fO2_offset = to_float(fO2_offset)
    
    comp = bulk.copy()

    if fO2_buffer is not None:
        if fO2_buffer != "NNO":
            if fO2_buffer != "FMQ":
                raise Warning("fO2 buffer specified is not an allowed input. This argument can only be 'FMQ' or 'NNO' \n if you want to offset from these buffers use the 'fO2_offset' argument.")

    if "MELTS" not in Model:
        if fO2_buffer == "FMQ":
            fO2_buffer = "qfm"
        if fO2_buffer == "NNO":
            fO2_buffer = "nno"
        _ensure_julia_ready()
            
    if H2O_Sat:
        comp['H2O_Liq'] = 20

    comp = comp_fix(Model=Model, comp=comp, Fe3Fet_Liq=Fe3Fet_init,
                    H2O_Liq=H2O_init, CO2_Liq=CO2_init)

    combined_results = {}
    if multi_processing:
        index_in = np.arange(len(P_bar))
        index_out = np.array([], dtype=int)

        while len(index_out) < len(index_in):
            Start = time.time()
            index = np.setdiff1d(index_in, index_out)
            groups = np.array_split(index, cores)
            non_empty_groups = [g for g in groups if g.size > 0]
            groups = non_empty_groups

            processes = []
            for i in range(len(groups)):
                q = Queue()
                p = Process(target = path_4_saturation_multi, args = (q, groups[i]),
                            kwargs = {'Model': Model, 'comp': comp,
                            'T_initial_C': T_initial_C, 'T_maxdrop_C': T_maxdrop_C,
                            'dt_C': dt_C, 'P_bar': P_bar, 'phases': phases,
                            'fO2_buffer': fO2_buffer, 'fO2_offset': fO2_offset})
                
                processes.append((p, q, groups[i]))
                p.start()

            if "MELTS" not in Model:
                for p, q, group in processes:
                    try:
                        if time.time() - Start > timeout + 10:
                            res = q.get(timeout = 10)
                        else:
                            res = q.get(timeout=timeout)
                    except:
                        if "MELTS" not in Model:
                            print(f"Timeout warning reached. Calculation P_bar = {P_bar[group[0]]} will not be returned. Try increasing the timeout.")
                        res = []
                        idx_chunks = np.array([group[0]], dtype = int)
                        mask = np.ones(len(P_bar), dtype=bool)
                        mask[idx_chunks] = False
                        P_bar = P_bar[mask]

                    p.join(timeout = 2)
                    p.terminate()

                    if len(res) > 0.0:
                        idx_chunks, results = res
                        combined_results.update(results)

                    index_out = np.hstack((index_out, idx_chunks))
            else:
                for p, q, group in processes:
                    while True:
                        try:
                            item = q.get(timeout=0.1)

                            run_index, step_i, data = item
                            combined_results.setdefault(run_index, {})[step_i] = data

                        except Empty:
                            time.sleep(0.001)
                            pass

                        # Timeout condition
                        if time.time() - Start > timeout:
                            print("Timeout — terminating process")
                            p.terminate()
                            break

                    # If process finished and queue empty → break
                        if not p.is_alive():
                            try:
                                # Drain remaining items
                                while True:
                                    item = q.get_nowait()
                                    run_index, step_i, data = item
                                    combined_results.setdefault(run_index, {})[step_i] = data
                            except Empty:
                                time.sleep(0.001)
                                pass
                            break

                index_out = np.array(list(combined_results.keys()))

    else:
        if "MELTS" in Model:
            from meltsdynamic import MELTSdynamic

            if Model is None or Model == "MELTSv1.0.2":
                melts = MELTSdynamic(1)
            elif Model == "pMELTS":
                melts = MELTSdynamic(2)
            elif Model == "MELTSv1.1.0":
                melts = MELTSdynamic(3)
            elif Model == "MELTSv1.2.0":
                melts = MELTSdynamic(4)
        else:
            from juliacall import Main as jl, convert as jlconvert
            env_dir = Path.home() / ".petthermotools_julia_env"
            jl_env_path = env_dir.as_posix()

            jl.seval(f"""
                import Pkg
                Pkg.activate("{jl_env_path}")
                """)

            jl.seval("using MAGEMinCalc")

            comp_julia = jl.seval("Dict")(comp) 

        if "MELTS" in Model:
            for i in range(len(P_bar)):
                results = {}
                Results = path_MELTS(P_bar = P_bar[i],
                        Model=Model,
                        comp=comp,
                        T_maxdrop_C=T_maxdrop_C,
                        dt_C=dt_C,
                        T_initial_C=T_initial_C,
                        find_liquidus=True,
                        fO2_buffer=fO2_buffer,
                        fO2_offset=fO2_offset,
                        fluid_sat=H2O_Sat,
                        Suppress=['rutile', 'tridymite'],
                        phases=phases,
                        melts = melts
                    )
                results[i] = Results
                combined_results.update(results)
        else:
            if fO2_offset is None:
                fO2_offset = 0.0

            for i in range(len(P_bar)):
                results = {}
                Results_df = jl.MAGEMinCalc.path(
                                        comp=comp_julia, dt_C=dt_C,
                                        P_bar=P_bar[i], T_min_C = T_maxdrop_C,            
                                        Model=Model, fo2_buffer=fO2_buffer, 
                                        fo2_offset=fO2_offset, find_liquidus=True,
                                        frac_xtal = False
                                        )
                
                Results = dict(Results_df)
                results[f"Run {i}"] = Results
                combined_results.update(results)
            
    results = combined_results

    if "MELTS" in Model:
        full_results = {}
        for run in results:
            if results[run]:
                if results[run][-1]:
                    del results[run][-1]
                res = results[run]
                new_results = compile_results(res)

                full_results[f"Run {run}"] = new_results

        Results = stich(Res=full_results, multi=True, Model=Model)
    else:
        Results = stich(Res=results, multi=True, Model=Model)

    ## determine the offset between the phases
    if len(phases)  == 3:
        arr = np.zeros((len(Results.keys()), 4))
        arr2 = np.zeros((len(Results.keys()), 4))
        columns = ['P_bar'] + phases + [phases[0] + ' - ' + phases[1], phases[0] + ' - ' + phases[2], phases[1] + ' - ' + phases[2], '3 Phase Saturation']
        for i, r in enumerate(Results.keys()):
            for idx, p in enumerate(phases):
                if p in Results[r]['mass_g'].keys():
                    arr[i, idx+1] = Results[r]['Conditions'].loc[Results[r]['mass_g'][p] > 0.0, 'T_C'].values[0]
                else:
                    arr[i, idx + 1] = np.nan

            arr[i,0] = Results[r]['Conditions']['P_bar'].loc[0]

        arr2[:, 0] = np.abs(arr[:,1] - arr[:,2])
        arr2[:, 1] = np.abs(arr[:,1] - arr[:,3])
        arr2[:, 2] = np.abs(arr[:,2] - arr[:,3])
        arr2[:, 3] = np.max(arr2[:,0:2], axis = 1)

        r_arr = np.hstack((arr,arr2))

        out = pd.DataFrame(data = r_arr, columns = columns)
        out = out.sort_values('P_bar').reset_index(drop=True)

    else:
        arr = np.zeros((len(Results.keys()),4))
        columns = ['P_bar'] + phases + [phases[0]+' - '+phases[1]]
        for i, r in enumerate(Results.keys()):
            for idx, p in enumerate(phases):
                if p in Results[r]['mass_g'].keys():
                    arr[i, idx+1] = Results[r]['Conditions'].loc[Results[r]['mass_g'][p] > 0.0, 'T_C'].values[0]
                else:
                    arr[i, idx + 1] = np.nan

            arr[i,0] = Results[r]['Conditions']['P_bar'].loc[0]

        arr[:,3] = np.abs(arr[:,1] - arr[:,2])

        out = pd.DataFrame(data = arr, columns = columns)
        out = out.sort_values('P_bar').reset_index(drop=True)

    if find_min:
        res = findmin(out = out, P_bar = P_bar, T_cut_C = T_cut_C)
        out = {'CurveMin': res, 'Output': out}
    else:
        out = {'Output': out}

    return out, Results

def findmin(out = None, P_bar = None, T_cut_C = None):
    """
    Finds the minimum temperature offset between mineral saturation temperatures as a function of pressure.

    Parameters
    ----------
    out : pd.DataFrame
        DataFrame containing phase saturation temperatures.
    P_bar : array-like
        Pressure values (bar) corresponding to the rows in `out`.
    T_cut_C : float
        Maximum acceptable temperature difference (°C) for co-saturation to be considered valid.

    Returns
    -------
    CurveMin : dict
        Dictionary containing:
            - 'P_min': Pressure of minimum saturation temperature difference.
            - 'Res_min': Minimum temperature difference.
            - 'y_new': Interpolated curve of temperature differences.
            - 'P_new': Pressure array for interpolated curve.
            - 'test': 'Pass' or 'Fail' depending on whether the result meets `T_cut_C`.
    """
    Res = out.copy()

    if '3 Phase Saturation' in list(Res.keys()):
        T_cut_old = T_cut_C
        T_cut_C = T_cut_C + np.nanmin(Res[4:])
        Minimum = list(Res.keys())[4:]
        CurveMin = {}
        for m in Minimum:
            if len(Res[m][~np.isnan(Res[m].values)]) > 2:
                y = Res[m][(~np.isnan(Res[m].values)) & (Res[m].values < T_cut_C)].values
                x = P_bar[(~np.isnan(Res[m].values)) & (Res[m].values < T_cut_C)]

                try:
                    y_new = interpolate.UnivariateSpline(x, y, k = 3)

                    P_new = np.linspace(P_bar[P_bar == np.nanmin(P_bar[(~np.isnan(Res[m].values)) & (Res[m].values < T_cut_C)])], 
                                        P_bar[P_bar == np.nanmax(P_bar[(~np.isnan(Res[m].values)) & (Res[m].values < T_cut_C)])], 200)

                    NewMin = np.nanmin(y_new(P_new))
                    P_min = P_new[np.where(y_new(P_new) == NewMin)][0]
                    if NewMin < T_cut_old:
                        Test = 'Pass'
                    else:
                        Test = 'Fail'

                    CurveMin[m] = {'P_min': P_min, 'Res_min': NewMin, 'y_new': y_new(P_new), 'P_new': P_new, 'test': Test}
                except:
                    try:
                        y_new = interpolate.UnivariateSpline(x, y, k = 2)

                        P_new = np.linspace(P_bar[P_bar == np.nanmin(P_bar[(~np.isnan(Res[m].values)) & (Res[m].values < T_cut_C)])], 
                                            P_bar[P_bar == np.nanmax(P_bar[(~np.isnan(Res[m].values)) & (Res[m].values < T_cut_C)])], 200)

                        NewMin = np.nanmin(y_new(P_new))
                        P_min = P_new[np.where(y_new(P_new) == NewMin)][0]
                        if NewMin < T_cut_old:
                            Test = 'Pass'
                        else:
                            Test = 'Fail'

                        CurveMin[m] = {'P_min': P_min, 'Res_min': NewMin, 'y_new': y_new(P_new), 'P_new': P_new, 'test': Test}
                    except:
                        CurveMin[m] = {'P_min': np.nan, 'Res_min': np.nan, 'y_new': np.nan, 'P_new': np.nan, 'test': 'Fail'}
            else:
                y_new = np.nan
                P_new = np.nan
                NewMin = np.nan
                P_min = np.nan
                Test = 'Fail'
                CurveMin[m] = {'P_min': P_min, 'Res_min': NewMin, 'y_new': y_new, 'P_new': P_new, 'test': Test}

    else:
        CurveMin = {}
        m = Res.keys()[3]
        if len(Res[m][~np.isnan(Res[m])]) > 2:
            y = Res[m][(~np.isnan(Res[m].values)) & (Res[m].values < T_cut_C*2)].values
            x = P_bar[(~np.isnan(Res[m].values)) & (Res[m].values < T_cut_C*2)]

            try:
                y_new = interpolate.UnivariateSpline(x, y, k = 3)
            except:
                y_new = interpolate.UnivariateSpline(x, y, k = 2)

            P_new = np.linspace(P_bar[P_bar == np.nanmin(P_bar[(~np.isnan(Res[m])) & (Res[m] < T_cut_C*2)])], 
                                    P_bar[P_bar == np.nanmax(P_bar[(~np.isnan(Res[m])) & (Res[m] < T_cut_C*2)])], 200)
            
            NewMin = np.nanmin(y_new(P_new))
            P_min = P_new[np.where(y_new(P_new) == NewMin)][0]
            if NewMin < T_cut_C:
                Test = 'Pass'
            else:
                Test = 'Fail'
        else:
            y_new = np.nan
            P_new = np.nan
            NewMin = np.nan
            P_min = np.nan
            Test = 'Fail'

        CurveMin[m] = {'P_min': P_min, 'Res_min': NewMin, 'y_new': y_new, 'P_new': P_new, 'test': Test}

    return CurveMin

def find_mineral_cosaturation(cores = None, Model = None, bulk = None, phases = None, P_bar = None, Fe3Fet_Liq = None, 
                              H2O_Liq = None, H2O_Sat = False, T_initial_C = None, dt_C = None, T_maxdrop_C = None, 
                              T_cut_C = None, find_range = None, find_min = None, fO2_buffer = None, fO2_offset = None):
    print("This function has been removed following update to v0.2.40. Please switch to using the mineral_cosaturation() function")

    return "This function has been removed following update to v0.2.40. Please switch to using the mineral_cosaturation() function"