from plotter import Plotter, U
import numpy as np
from functools import partial
import os
import pandas as pd
import tables
from turbox import *
from ismfeed import get_sinks, compute_sfr
from scipy.stats import linregress
# from astropy.modeling import models, fitting
# from astropy.stats import gaussian_fwhm_to_sigma
# from models import Rabatin


## Power spectrum

def do_pspec(pp):
    try: 
        return pp.pspec(magnetic=False)
    except:
        pp.logger.error("Failed to run pspec")



def do_all_pspec(pl):

    pl.map(do_pspec)
    pl.turbox["slope_pspec_logrho"] = np.zeros(len(pl.turbox))
    pspec = {}
    k_dict = {}

    for run in pl.runs:
        s = 0
        all_pspec = []
        sel_nums = pl.get_snap_list(select={"runs":run})
        for pp in sel_nums:
            logkmin, logkmax = span_resolution[256]
            k, power = get_pspec(pp, "logrho")
            logk  = np.log10(k)
            mask = (logk >= logkmin) & (logk < logkmax)
            all_pspec.append(power)

        all_pspec = np.array(all_pspec)    
        k_dict[run] = k
        pspec[run] = np.sum(all_pspec, axis=0) / len(sel_nums)
        results = linregress(logk[mask], np.log10(pspec[run][mask]))
        pl.turbox["slope_pspec_logrho"].loc[run] = results.slope
    pl.turbox["pspec_logrho"] = pspec.values()
    pl.turbox["pspec_k"] = k_dict.values()
    

def do_all_pdf(pl):
    pl.map(apply_rule_pdf) 
    pl.study.time_rho_pdf(overwrite=True, skip_dep=False)
    pl.study.avg_time_rho_pdf(overwrite=True, skip_dep=True)

    pdf_data = pl.study.get_value("/comp/avg_time_rho_pdf")["median"]
    q16 = pl.study.get_value("/comp/avg_time_rho_pdf")["q16"]
    q84 = pl.study.get_value("/comp/avg_time_rho_pdf")["q84"]

    pl.turbox["pdf_logrho_x"] = list(pdf_data[:, 1])
    pl.turbox["pdf_logrho_y"] = list(pdf_data[:, 0])
    pl.turbox["pdf_logrho_ymin"] = list(q84[:, 0])
    pl.turbox["pdf_logrho_ymax"] = list(q16[:, 0])
    
    logrho_mean = {}
    logrho_var = {}
    for run in pl.runs:
        logrho = pl.turbox["pdf_logrho_x"][run]
        logrho_pdf = pl.turbox["pdf_logrho_y"][run]
        dV = logrho_pdf * np.diff(logrho)[0]
        logrho_mean[run] = np.sum(logrho * dV)
        logrho_var[run] = np.sum((logrho - logrho_mean[run])**2 * dV)
    
    pl.turbox["mean_logrho_V"] = list(logrho_mean.values())
    pl.turbox["S_logrho_V"] = list(logrho_var.values())

  
def add_sfr(pl, pl_grav, dep=False):
    
    if dep:
        pl_grav.study.nml("turb_params/comp_frac",overwrite=True, overwrite_dep=True)
        pl_grav.study.nml("turb_params/turb_rms", overwrite=True)
        pl_grav.study.sinks_from_log()
        pl_grav.study.coarse_step_from_log()
    get_sinks(pl_grav.study, stellar=False, sk=False)
        
    if dep or not hasattr(pl_grav.study, "sfr"):
        sfr, sfr_err = compute_sfr(pl_grav.study, target_start=0.005, target_end=0.2)
    else:
        sfr, sfr_err = pl_grav.study.sfr, pl_grav.study.sfr_err
        
        
    pl.turbox["sfr"] = np.zeros(len(pl.runs))
    pl.turbox["sfr_err"] =  np.zeros(len(pl.runs))

    for i, run in enumerate(pl_grav.runs):
        pl.turbox["sfr"][run] = sfr[i]
        pl.turbox["sfr_err"][run] = sfr_err[i]



def do_pspec_pdf_sfr(pl, pl_grav, redo=False):
    
    build_turbox_data(pl, redo=redo)
    do_all_pspec(pl)
    do_all_pdf(pl)
    add_sfr(pl, pl_grav, dep=redo)
    pl.turbox.to_csv(f"turbox_{pl.params.out.tag}.csv")
    pl.turbox.to_pickle(f"turbox_{pl.params.out.tag}.pickle")


if __name__ == "__main__":
        
    in_dir_grav = os.path.expanduser("~nbrucy/simus/turbox/grav")
    out_dir_grav =  os.path.expanduser("~nbrucy/visus/turbox/grav")

    in_dir_nograv = os.path.expanduser("~nbrucy/simus/turbox/nograv")
    out_dir_nograv =  os.path.expanduser("~nbrucy/visus/turbox/nograv")

    cs0 = 0.28834810480560674

    pl_nograv = Plotter(
        in_dir_nograv, 
        path_out=out_dir_nograv, 
        unit_time="turb_params/turb_T",
        #time=4,
        time_min=2,
        time_max=8,
        tag="nograv",
        filter_name="n1.5*",
        params="turbox_params.yml",
        )


    pl_nograv_L200 = Plotter(
        in_dir_nograv + '/L200', 
        path_out=out_dir_nograv  + '/L200', 
        unit_time="turb_params/turb_T",
        #time=4,
        time_min=2,
        time_max=8,
        tag="nograv_L200",
        filter_name="*comp0.5*",
        params="turbox_params.yml",
        )


    pl_grav = Plotter(
        in_dir_grav, 
        path_out=out_dir_grav, 
        unit_time="turb_params/turb_T",
        time=2,
        tag="grav",
        filter_name="n1.5*",
        params="turbox_params.yml",
        )


    pl_grav_L200 = Plotter(
        in_dir_grav + '/L200', 
        path_out=out_dir_grav  + '/L200', 
        unit_time="turb_params/turb_T",
        time=2,
        tag="grav_L200",
        filter_name="*comp0.5*",
        params="turbox_params.yml",
        )

    do_pspec_pdf_sfr(pl_nograv, pl_grav, False)
    do_pspec_pdf_sfr(pl_nograv_L200, pl_grav_L200, False)