import os
import pandas as pd
import yaml
from assets.constants import config_overrides
from typing import (Union,
                    Tuple)
from assets.commons import validate_parameter


def get_poc_results(line_fit_filename: str) -> pd.DataFrame:
    df_class = pd.read_csv(os.path.join('assets', 'data', line_fit_filename))
    with open(os.path.join('prs', 'output', 'run_type', 'constant_abundance_p15_q05', 'volume_density_results.yml'),
              'r') as infile:
        fit_results = yaml.load(infile, Loader=yaml.FullLoader)
    df_results = pd.DataFrame.from_dict(fit_results, orient='index')
    df_full = (df_results.merge(df_class, left_index=True, right_on='source_name', how='inner', validate='1:1')
               .dropna(subset='best_fit'))
    df_full['class_phys'] = 'HII'
    df_full.loc[df_full['class'] == '70w', 'class_phys'] = 'Quiescent'
    df_full.loc[df_full['class'] == 'IRw', 'class_phys'] = 'Protostellar'
    df_full.loc[df_full['class'] == 'IRb', 'class_phys'] = 'MYSOs'
    return df_full


def get_profile_input(
        model_root_folder: str,
        species: Union[None, str] = None) -> Tuple[dict, dict]:
    overrides = config_overrides[model_root_folder.split(os.path.sep)[-1]]
    _species = validate_parameter(species, default='e-ch3oh')
    try:
        overrides_grid = overrides['grid']
    except KeyError:
        overrides_grid = {}
    try:
        overrides_lines = overrides['lines']
    except KeyError:
        overrides_lines = {}
    overrides_abundance = {}

    try:
        overrides_abundance['molecular_abundance'] = overrides_lines['molecular_abundances'][_species]
    except KeyError:
        pass

    for key in ['abundance_jump', 'threshold']:
        try:
            overrides_abundance[key] = overrides_lines['hot_core_specs'][_species][key]
        except KeyError:
            pass

    return overrides_grid, overrides_abundance

