#!/usr/bin/env python
# coding: utf-8

"""\
Code for error parametrization regresison functions:
- CO2 satellite obervations

Equations are from LOGOFLUX-1 report
"""

import numpy as np


def valid_range(x, xmin, xmax):
    return (xmin <= x) & (x <= xmax)


def epf(sza, albedo_nir, albedo_swir, aod=0.0, cod=0.0, cth=10.0,
        error_type='random', species='CO2'):
    """\

    parameter   | description                       | valid range
    ------------+-----------------------------------+------------------------
    sza         | solar zenith angle (degree)       | 0 - 80 degree
    albedo_nir  | albedo NIR band                   | 0.03 - 0.7
    albedo_swir | albedo SWIR band                  | 0.03 - 0.7
    aod         | aerosol optical depth at 550 nm   | 0 - 0.6
    cod         | cirrus optical depth (NIR)        | 0 - 0.6
    cth         | cirrus top height (km)            | 2 - 20 km
    ------------+-----------------------------------+------------------------
    """
    # find shape
    for x in [sza, albedo_nir, albedo_swir, aod, cod, cth]:
        if np.ndim(x) != 0:
            shape = x.shape
            break

    # either raise if out of range or mask
    valids = (
        valid_range(sza,         0.0, 80.0) &
        valid_range(albedo_nir,  0.03, 0.7) &
        valid_range(albedo_swir, 0.03, 0.7) &
        valid_range(aod,         0.0,  0.6) &
        valid_range(cod,         0.0,  0.6) &
        valid_range(cth,         2.0, 20.0)
    )

    inc_sza = np.cos(np.deg2rad(84.0)) / np.cos(np.deg2rad(sza + 9.0)) * sza / 75.0
    inc_alb = (1.01 / (albedo_swir + 0.01) - 1.0) * 0.01

    x = np.array([
        np.ones(shape),
        sza - np.full(shape, 50.0),
        albedo_nir - np.full(shape, 0.10),
        albedo_swir - np.full(shape, 0.10),
        aod - np.full(shape, 0.2),
        cod - np.full(shape, 0.05),
        cth - np.full(shape, 10.0),
        aod * inc_sza * inc_alb
    ])

    # CO2/CH4 relative error
    if error_type == 'random':
        c = {
            'CO2': np.array([1.29121, -0.00220, -0.07275, -1.97335,  -1.00620,  0.34586, -0.00150, 38.86071]),
            'CH4': np.array([9.47532, -0.02678, -7.84754, -11.52154, -7.88635, 15.80352,  0.01702, 275.25335])
        }[species]

    # CO2/CH4 sys. errors
    if error_type == 'sys':
        c = {
            'CO2': np.array([0.30786, 0.01191, 1.97716, -1.55154, -1.40282, 0.70400, -0.03020, -73.46152]),
            'CH4': np.array([1.41086, 0.05545, 9.17830, -6.77504, -9.15046, 5.52796, -0.48784, -335.04555])
        }[species]


    # calculate error
    errors = np.zeros(shape)
    for i in range(c.size):
        errors += c[i] * x[i]

    # "post-processing"
    if error_type == 'random':
        errors[errors < 0.7] = 0.7

    if error_type == 'sys':
        errors -= sza / 70.0 - 0.2

    return np.ma.array(errors, mask=~valids)






