#!/usr/bin/env python

"""fetch_simulated_data.py: fetch simulated (WACCM-X) data from netcdf files."""

__author__ = "<NAME>."
__copyright__ = "Copyright 2020, SuperDARN@VT"
__credits__ = []
__license__ = "MIT"
__version__ = "1.0"
__maintainer__ = "<NAME>."
__email__ = "<EMAIL>"
__status__ = "Research"

import sys
sys.path.extend(["code/", "code/rt/", "code/sd/"])

import os
import datetime as dt
import numpy as np
import netCDF4 as nc

from scipy import constants as const
from scipy.interpolate import interp1d, interp2d
import pickle
from scipy import array


pconst = {
    "boltz" : 1.38066e-23,       # Boltzmann constant  in Jule K^-1
    "h"     : 6.626e-34  ,       # Planks constant  in ergs s
    "c"     : 2.9979e08  ,       # in m s^-1
    "avo"   : 6.023e23   ,       # avogadro's number           
    "Re"    : 6371.e3    ,
    "amu"   : 1.6605e-27 ,
    "q_e"   : 1.602e-19  ,       # Electron charge in C
    "m_e"   : 9.109e-31  ,       # Electron mass in kg
    "g"     : 9.81       ,       # Gravitational acceleration on the surface of the Earth
    "eps0"  : 1e-9/(36*np.pi),
    "R"     : 8.31       ,       # J mol^-1 K^-1
}

def extrap1d(x,y,kind="linear"):
    """ This method is used to extrapolate 1D paramteres """
    interpolator = interp1d(x,y,kind=kind)
    xs = interpolator.x
    ys = interpolator.y
    def pointwise(x):
        if x < xs[0]: return ys[0]+(x-xs[0])*(ys[1]-ys[0])/(xs[1]-xs[0])
        elif x > xs[-1]: return ys[-1]+(x-xs[-1])*(ys[-1]-ys[-2])/(xs[-1]-xs[-2])
        else: return interpolator(x)
    def ufunclike(xs):
        return array(list(map(pointwise, array(xs))))
    return ufunclike

def _get_ij_(lats, lons, lat, lon):
    """ Get (lat, lon) index """
    _ij_ = (np.argmin(np.abs(lats-lat)), np.argmin(np.abs(lons-lon)))
    return _ij_

def _intp_heights_(h, hx, param, scale="linear", kind="linear"):
    if scale == "linear": pnew = extrap1d(h, param, kind=kind)(hx)
    if scale == "log": pnew = 10**extrap1d(h, np.log10(param), kind=kind)(hx)
    return pnew

def _intp_latlon_(lats, lons, latx, lonx, param, scale="linear", kind="cubic"):
    if scale == "linear": pnew = interp2d(lats, lons, param.T, kind=kind)(latx, lonx).T
    if scale == "log": pnew = 10**interp2d(lats, lons, np.log10(param.T), kind=kind)(latx, lonx).T
    return pnew

def _intp_(h, lats, lons, param, hd=[50,600,1], dlat=0.5, dlon=1, scale="log", kind="cubic", v=False):
    print(" Defined scale - ", scale)
    lonx = np.arange(np.min(lons), np.max(lons), dlon)
    latx = np.arange(np.min(lats), np.max(lats), dlat)
    param_intm = np.zeros((len(h), len(latx), len(lonx)))
    h_intm = np.zeros((len(h), len(latx), len(lonx)))
    for k,_ in enumerate(h):
        param_intm[k, :, :] = _intp_latlon_(lats, lons, latx, lonx, param[k, :, :], scale, kind)
        h_intm[k, :, :] = _intp_latlon_(lats, lons, latx, lonx, h[k, :, :], scale, kind)
    if v: print("\tLatlon convt.")
    hx = np.arange(hd[0],hd[1],hd[2])
    pnew = np.zeros((len(hx), len(latx), len(lonx)))
    for i,_ in enumerate(latx):
        for j,_ in enumerate(lonx):
            if scale == "linear": pnew[:,i,j] = extrap1d(h_intm[:, i, j], param_intm[:, i, j])(hx)
            if scale == "log": pnew[:,i,j] = 10**(extrap1d(h_intm[:, i, j], np.log10(param_intm[:, i, j]))(hx))
            #pnew[:,i,j] = 10**(extrap1d(h_intm[:, i, j], np.log10(param_intm[:, i, j]))(hx))
    if v: print("\tHeight convt.")
    return pnew, hx, latx, lonx

def transform_ne(T, e, lev, unit="cm"):
    """
    Transform electron density from mol/mol to /cc or /cm
    Formula
    -------
    P = nKT, K = Boltzmann constant in SI
    lev: pressure in hPa (1hPa=100Pa)
    T: neutral temperature
    e: electron density in mol/mol
    
    T in K, P in Pa, n is electron density in /cubic meter or /cc
    """
    P = lev * 1e2
    u = 1 if unit == "cm" else 1e-6
    ne = np.zeros_like(e)
    for i, p in enumerate(P):
        na = p * pconst["avo"] / (pconst["R"] * T[:,i,:,:])
        n = u * na * e[:,i,:,:]
        ne[:,i,:,:] = n
    print(" Maxima and minima in Ne - (in 1e12)", np.round(ne.max()/1e12, 3), np.round(ne.min()/1e12, 2))
    return ne

def extract_electron_density(ev, lat, lon, dates, params, fname):
    ds = nc.Dataset(fname)
    _dict_ = {}
    T = [ev + dt.timedelta(seconds=float(d)) for d in ds.variables["datesec"][:]]
    ti, tj = T.index(dates[0]), T.index(dates[1])
    _dict_["time"] = T[ti:tj]
    latx = ds.variables["lat"][:]
    lonx = ds.variables["lon"][:]
    lonx = np.where(lonx > 180, lonx-360, lonx)
    i, j = _get_ij_(latx, lonx, lat, lon)
    _dict_["lev"] = ds.variables["lev"][:]
    _dict_["Z3GM"] = ds.variables["Z3GM"][ti:tj+1,:,i,j]
    T, e = self.nc.variables["T"][:], self.nc.variables["e"][:]
    ne = transform_ne(T, e, _dict_["lev"])
    _dict_["ne"] = ne[ti:tj+1,:,i,j]
    return _dict_

def extract_parameters(ev, lat, lon, dates, params, fnames, keys, mult=1., scale="log", 
                       kind="linear"):
    _dict_ = {}
    for key, fname in zip(keys, fnames):
        _dict_[key] = {}
        ds = nc.Dataset(fname)
        T = [ev + dt.timedelta(seconds=float(d)) for d in ds.variables["datesec"][:]]
        ti, tj = T.index(dates[0]), T.index(dates[1])
        _dict_[key]["T"] = T[ti:tj+1]
        latx = ds.variables["lat"][:]
        lonx = ds.variables["lon"][:]
        lonx = np.where(lonx > 180, lonx-360, lonx)
        i, j = _get_ij_(latx, lonx, lat, lon)
        hx = np.arange(50,600,1)
        _dict_[key]["H"] = hx
        Z3 = ds.variables["Z3GM"][ti:tj+1,:,i,j]/1e3
        for p in params:    
            ox = ds.variables[p][ti:tj+1,:,i,j]
            px = np.zeros((ox.shape[0], len(hx)))
            for i in range(ox.shape[0]):
                px[i,:] = _intp_heights_(Z3[i,:], hx, ox[i,:], scale=scale, kind=kind)
            _dict_[key][p] = px * mult
    return _dict_

