from oggm import entity_task
from oggm.core.flowline import run_from_climate_data
from oggm.core.massbalance import MonthlyTIModel
from functools import partial

# Module logger
import logging
log = logging.getLogger(__name__)

MonthlyTIModel_check = partial(MonthlyTIModel, check_calib_params=False)

@entity_task(log, writes=['mb_calib'])
def perturbate_mb_params(gdir, perturbation=None, reset_default=False):
    """Replaces existing MB params with perturbed ones.

    Assumes that the parameter perturbations are additive,
    except for precip which is multiplicative.

    The default original params are kept in the dictionary for re-use.

    Parameters
    ----------
    parameters : dict
        the parameters to change and their value
    reset_default : bool
        reset the parameters to their original value
    """

    df = gdir.read_json('mb_calib')

    # Save original params if not there
    if 'bias_orig' not in df:
        for k in ['bias', 'melt_f', 'prcp_fac', 'temp_bias']:
            df[k + '_orig'] = df[k]

    if reset_default:
        for k in ['bias', 'melt_f', 'prcp_fac', 'temp_bias']:
            df[k] = df[k + '_orig']
        gdir.write_json(df, 'mb_calib')
        return df

    for k, v in perturbation.items():
        if k == 'prcp_fac':
            df[k] = df[k + '_orig'] * v
        else:
            df[k] = df[k + '_orig'] + v

    gdir.write_json(df, 'mb_calib')
    return df


@entity_task(log)
def run_from_climate_data_glen_a_fac(gdir, glen_a_fac=None, **kwargs):
    """Adds magic to the simulations

    Parameters
    ----------
    glen_a_fac : float 
        multiplicative 
    """
    
    if glen_a_fac is None:
        glen_a_fac = 1
        
    diag = gdir.get_diagnostics()
    glen_a = diag.get('inversion_glen_a') * glen_a_fac
        
    return run_from_climate_data(gdir, glen_a=glen_a, mb_model_class=MonthlyTIModel_check, **kwargs)