from scipy.integrate import solve_ivp
from plotter import U
import select_snapshot
import numpy as np
import pandas as pd

ssfr_base = 2.5e-10


def get_sinks(self, overwrite=False, stellar=True, convert_units=False):
    
    self.sinks_from_log(overwrite=overwrite)
    self.coarse_step_from_log(overwrite=overwrite)
    self.sinks = self.get_value("/series/sinks_from_log")
    
    if stellar:
        self.stellar_from_log(overwrite=overwrite)
        self.stellar = self.get_value("/dataset/stellar_from_log")
        self.stellar["sn_time"] = {}
        for run in self.runs:
            self.stellar["sn_time"][run] =  (self.stellar["time"][run]  + self.stellar["lifetime"][run] )*self.info["unit_time"].express(U.year)
            ind_sort = self.stellar["sn_time"][run].argsort()
            for key in self.stellar:
                self.stellar[key][run] = self.stellar[key][run][ind_sort]
        self.stellar["cum_mass"] = {}
        sn = {}
        sn["time"] = {}
        sn["cum_mass"] = {}
    
    self.sinks["cum_mass"] = {}

    for run in self.runs:
        
        if convert_units:
            self.sinks["time"][run] *= self.info["unit_time"].express(U.year)
        
        self.sinks["cum_mass"][run] = self.sinks["mass_sink"][run].copy()
        
        if stellar:
            self.stellar["cum_mass"][run] = np.cumsum(self.stellar["mass"][run]) * self.info["unit_mass"].express(U.Msun)
            sn["time"][run], idx, count = np.unique(self.stellar["sn_time"][run], return_index=True, return_counts=True) 
            idx += count - 1
            sn["cum_mass"][run] = self.stellar["cum_mass"][run][idx]
            ind_sn = np.searchsorted(self.sinks["time"][run], sn["time"][run])
            for i, indi in enumerate(ind_sn[ind_sn < ind_sn[-1]]):
                self.sinks["cum_mass"][run][indi:ind_sn[i+1]] +=  sn["cum_mass"][run][i]
            

    time_gas, total_mass, mass_gas = self.total_mass()
    
    sk_ssfr = {}
    for run in self.runs:
        time_gas[run] *= self.info["unit_time"].express(U.Myr)
        sk_ssfr[run] = ssfr_base * 1e6 * (mass_gas[run]/1e6)**1.4

    def dermass(t, sm, run, fact_sfr=1):
        ind = min(np.searchsorted(time_gas[run], t), sk_ssfr[run].size - 1)
        return sk_ssfr[run][ind]*fact_sfr

    tspan = np.linspace(0, 200, 100)

    self.time_esm = {}
    self.expected_stellar_mass = {}
    self.expected_stellar_mass_max = {}
    self.expected_stellar_mass_min = {}

    for run in self.runs:
        sol = solve_ivp(dermass, (tspan[0], tspan[-1]), [0], t_eval=tspan, args=[run])
        self.time_esm[run], self.expected_stellar_mass[run] = sol["t"], sol["y"][0]
        sol_max = solve_ivp(dermass, (tspan[0], tspan[-1]), [0], t_eval=tspan, args=[run, 3])
        self.expected_stellar_mass_max[run] = sol_max["y"][0]
        sol_min = solve_ivp(dermass, (tspan[0], tspan[-1]), [0], t_eval=tspan, args=[run, 1./3])
        self.expected_stellar_mass_min[run] = sol_min["y"][0]
 
def compute_sfr(self, target_start=0.05, target_end=0.3):
    mass = self.sinks["cum_mass"]
    time = self.sinks["time"]
    sfr = {}
    sfr_err = {}
    tend = {}
    tstart = {}
    for run in self.runs:
        tend[run] = select_runs.time_mcons(self, run, target=target_end) * 1e6
        tstart[run] = select_runs.time_mcons(self, run, target=target_start) * 1e6
        if tend[run] < tstart[run]:
            tend[run] = time[run][-1]


    for run in self.runs:
        idx1 = time[run].searchsorted(tstart[run])
        idx2 = min(time[run].searchsorted(tend[run]), len(time[run]) - 1)

        sfr[run] = (mass[run][idx2] - mass[run][idx1]) /  (time[run][idx2] - time[run][idx1])
        sfr_other = []
        for i in range(idx1, idx2-10):
            sfr_other.append((mass[run][i+10:idx2+1] - mass[run][i]) /  (time[run][i+10:idx2+1] - time[run][i]))

        sfr_err[run] = np.std(np.concatenate(sfr_other))
    self.sfr = np.array(list(sfr.values()))
    self.sfr_err = np.array(list(sfr_err.values()))
    return self.sfr, self.sfr_err


def compute_power(self, target_start=0.05, target_end=0.3):
    mass = self.sinks["cum_mass"]
    time = self.sinks["time"]
    sfr = {}
    sfr_err = {}
    tend = {}
    tstart = {}
    for run in self.runs:
        tend[run] = select_runs.time_mcons(self, run, target=target_end) * 1e6
        tstart[run] = select_runs.time_mcons(self, run, target=target_start) * 1e6
        if tend[run] < tstart[run]:
            tend[run] = time[run][-1]


    for run in self.runs:
        idx1 = time[run].searchsorted(tstart[run])
        idx2 = min(time[run].searchsorted(tend[run]), len(time[run]) - 1)

        sfr[run] = (mass[run][idx2] - mass[run][idx1]) /  (time[run][idx2] - time[run][idx1])
        sfr_other = []
        for i in range(idx1, idx2-10):
            sfr_other.append((mass[run][i+10:idx2+1] - mass[run][i]) /  (time[run][i+10:idx2+1] - time[run][i]))

        sfr_err[run] = np.std(np.concatenate(sfr_other))
    self.sfr = np.array(list(sfr.values()))
    self.sfr_err = np.array(list(sfr_err.values()))
    return self.sfr, self.sfr_err