from __future__ import division

import numpy as np
from scipy.optimize import root
from .gas_cycle.fair1 import _iirf_simple, _iirf_interp
from .forcing.ghg import co2_log
from .defaults import carbon, thermal
from .constants import molwt
from .constants.general import ppm_gtc
from .temperature.millar import forcing_to_temperature, calculate_q


def infer_emissions(e1, c1_prescribed, carbon_boxes0, tau_new, a, c_pi):
    """Matches prescribed concentrations to forward-calculated concentrations.

    Inputs:
        e1            : emissions in timestep t, GtC
        c1_prescribed : CO2 concentrations in timestep t, ppmv
        carbon_boxes0 : carbon stored in each atmospheric reservoir at timestep
                        t-1 (GtC)
        tau_new       : decay time constants of CO2 (yr)
        a             : partition coefficient of carbon boxes
        c_pi          : pre-industrial concentration of CO2, ppmv
    """
    
    c1_calculated = np.sum(carbon_boxes0*np.exp(-1.0/tau_new) + a*e1 / ppm_gtc) + c_pi
    return c1_calculated-c1_prescribed


def inverse_carbon_cycle(c1, c_acc0, temp, r0, rc, rt, iirf_max, time_scale_sf,
                         a, tau, iirf_h, carbon_boxes0, c_pi, c0, e0):
    """Calculates CO2 emissions from concentrations.
    
    Inputs:
        c1            : concentration of CO2 in timestep t, ppmv
        c_acc0        : cumulative airborne carbon anomaly (GtC) since
                        pre-industrial, timestep t-1
        temp          : temperature anomaly above pre-industrial (K)
        r0            : pre-industrial time-integrated airborne fraction (yr)
        rc            : sensitivity of time-integrated airborne fraction to 
                        airborne carbon (yr/GtC)
        rt            : sensitivity of time-integrated airborne fraction to
                        temperature (yr/K)
        iirf_max      : maximum value of time-integrated airborne fraction (yr)
        time_scale_sf0: initial guess of alpha scaling factor
        a             : partition coefficient of carbon boxes
        tau           : present-day decay time constants of CO2 (yr)
        iirf_h        : time horizon for time-integrated airborne fraction (yr)
        carbon_boxes0 : carbon stored in each atmospheric reservoir at timestep
                        t-1 (GtC)
        c_pi          : pre-industrial concentration of CO2, ppmv
        c0            : concentration of CO2 in timestep t-1, ppmv
        e0            : emissions of CO2 in timestep t, GtC

    Outputs:
        e1            : emissions of CO2 in timestep t, GtC
        c_acc1        : cumulative airborne carbon anomaly (GtC) since
                        pre-industrial, timestep t
        carbon_boxes1 : carbon stored in each atmospheric reservoir at timestep
                        t (GtC)
        time_scale_sf : scale factor for CO2 decay constants
    """
     
    iirf = _iirf_simple(c_acc0, temp, r0, rc, rt, iirf_max)
    time_scale_sf = root(_iirf_interp, time_scale_sf,
        args=(a, tau, iirf_h, iirf))['x']
    tau_new = tau * time_scale_sf
    e1 = root(infer_emissions, e0, args=(c1, carbon_boxes0, tau_new, a, c_pi))['x']
    c_acc1 = c_acc0 + 0.5*(e1 + e0) - (c1 - c0)*ppm_gtc
    carbon_boxes1 = carbon_boxes0*np.exp(-1.0/tau_new) + a*e1 / ppm_gtc