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

# In[1]:


import numpy as np
import scipy.interpolate as spint
from scipy import optimize
import matplotlib.pyplot as plt
import pandas as pd
from scipy.spatial import distance

def deltaAZ(em, inc, ph_ang):
    c_em = np.cos((np.pi/180.)*em)
    s_em = np.sin((np.pi/180.)*em)
    
    c_inc = np.cos((np.pi/180.)*inc)
    s_inc = np.sin((np.pi/180.)*inc)
    
    c_phAng = np.cos((np.pi/180.)*ph_ang)
    
    dAZ = np.arccos((c_phAng - c_em*c_inc)/(s_em*s_inc))
    return(dAZ * 180./np.pi)

# interp documentation:
# wv, spfun = titanfit.interp(spapath,info)
# 
# creates interpolator function utilizing generated spectra from pydisort. function generated by regular grid interpolator function within scipy. 
# 
# parameters:
# 
# spapath = .npy file array of generated spectra using pydisort.
# info = tuple of: incidence, emission, delta azimuth, albedo, haze, and ch4 values.
# wvpath = .npy file array of wavelengths.
# 
# returns:
# 
# wv = array of wavelengths.
# spfun = interpolator function using regular grid interpolator. by default, bounds_error = False and fill_value = None.

# In[ ]:


def interp(spapath, wvpath, info):
    spa = np.load(spapath)
    wv = np.load(wvpath)
    iI = np.array(info[0]) # Incidence angle
    iE = np.array(info[1]) # emission angle
    iP = np.array(info[2]) # delta azimuth between inc and em
    iA = np.array(info[3]) # Albedo values
    iH = np.array(info[4]) # Haze opacity, relative to nominal
    iC = np.array(info[5]) # CH4 column, relative to nominal
    points=(iI,iE,iP,iA,iH,iC,wv)
    spfun = spint.RegularGridInterpolator(points, spa, bounds_error=False, fill_value=None)
    return wv, spfun


# grid documentation:
# grid = titanfit.grid(sequencepath,cube, latmin, latmax, lonmin, lonmax)
# 
# creates a reduced pandas dataframe of a specific cube from a given flyby. intended to speed up indexing when pulling spectra. currently, this step is unskippable. grids of each cube must be generated. 
# 
# parameters:
# 
# sequencepath = direct path to flyby/sequence .csv file.
# cube = full .cub cube name (ie 'V1643411254_1.ir_rc20final.cub'). must be within flyby/sequence.
# latmin, latmax = minimum latitude and maximum latitude range.
# lonmin, lonmax = minimum longitude and maximum longitude range.
# 
# returns:
# grid = pandas dataframe of cube over given longitude latitude range. 

# In[ ]:


def grid(sequencepath,cube, latmin, latmax, lonmin, lonmax):
    def gridlocator(cube, latmin, latmax, lonmin, lonmax):
        data = pd.read_csv(sequencepath)
        data2 = data[data['name']==f'{cube}']
        data3 = data2[data2['latitude'] > latmin]
        data4 = data3[data3['latitude'] < latmax]
        data5 = data4[data4['longitude'] > lonmin]
        data6 = data5[data5['longitude'] < lonmax]
        return data6
    return gridlocator(cube, latmin, latmax, lonmin, lonmax)


# vimsspec documentation:
# geo, wavg = vimsspec(cube, latitude, longitude, grid)
# 
# pulls the four nearest vims spectra from a given latitude, longitude point from a given cube and generates an average spectra by weighing by euclidean distance. 
# 
# parameters:
# 
# cube = full name of vims cube to pull from (ie 'V1643411254_1.ir_rc20final.cub'). must be within grid.
# latitude = latitude point to generate spectra nearest to.
# longitude = longitude point to generate spectra nearest to.
# grid = reduced pandas dataframe of a specific cube generated utilizing titanfit.grid. whole flybys cannot currently be utilized.
# 
# returns:
# 
# geo = viewing geometry dictionary of generated spectra created utilizing weighted average of the 4 spectras given viewing geometry. 
# wavg = average spectra generated utilizing weighted euclidean distances of the four nearest vims spectra.

# In[ ]:


def vimsspec(cube, latitude, longitude, grid):
    data2lonlat = grid[['latitude', 'longitude']]
    data2np = data2lonlat.to_numpy()
    idx = len(data2np)
    lst1 = []
    for a in range(idx):
        x = (data2np[a])
        y = (latitude, longitude)
        lst1.append(distance.euclidean(x,y))
    data2lonlat['euclidean'] = lst1

    result = data2lonlat['euclidean'].argsort()[:4]
    idx = []
    for a in result.values:
        idx.append(data2lonlat.iloc[a])
    idxdf = pd.DataFrame(idx)
    check = np.array(idxdf['euclidean'])
    spectra = []
    for a in idxdf.index:
        spectra.append(grid.loc[a])
    for a in spectra:
        avgin = np.average(a['incidence'])
    for a in spectra:
        avgem = np.average(a['emission'])
    for a in spectra:
        avgph = np.average(a['phase'])
    spectra2 = []
    # finds the spectra from the cube dictionary, then converts to an array.
    for a in spectra:
        teststr = str(a['spectra']).strip('{}')
        teststr2 = teststr.strip("''")
        spectra2.append(eval(teststr2))
    geo = {
      "incidence": avgin,
      "emission": avgem,
      "phase": avgph
    }
    totalweigh = check[0]+check[1]+check[2]+check[3]
    finalvals = []
    val = check[0]
    #speccheck = int(spectra2[0])
    avgspec = np.array(spectra2)
    wavg = (avgspec[0]*(check[0]/totalweigh))+(avgspec[1]*(check[1]/totalweigh))+(avgspec[2]*(check[2]/totalweigh))+(avgspec[3]*(check[3]/totalweigh))
    #avgspectogether = finalvals[0]+finalvals[1]+finalvals[2]+finalvals[3]
    #avgspeccheck = avgspectogether*(1/check[0]+check[1]+check[2]+check[3])
    # current weighing. 
    #avgspec2 = (avgspec[0]+avgspec[1]+avgspec[2]+avgspec[3])/4
    return geo, wavg


# coordgen documentation:
# coords = coordgen(lonmin, lonmax, latmin, latmax, resolution=625)
# generates a coordinate grid over a given longitude latitude range with a given resolution.
# 
# parameters:
# 
# lonmin, lonmax = the minimum and maximum longitude to generate grid.
# latmin, latmax = the minimum and maximum latitude to generate grid.
# resolution (optional): resolution to generated coordinate grid over. default is 625. (results in 625 points).
# 
# returns:
# 
# coords = list of coordinate pairs over longitude latitude range with given resolution. 

# In[ ]:


def coordgen(lonmin, lonmax, latmin, latmax, resolution=625):
    sq = np.sqrt(resolution)
    longitudes = np.linspace(lonmin, lonmax, int(sq))
    latitudes = np.linspace(latmin, latmax, int(sq))
    points = np.array([longitudes, latitudes])
    coords = []
    for a in points[0]:
        for b in points[1]:
            coords.append((a,b))
    return coords


# geospec documentation:
# compspec, vgeo = geospec(cubes, coords, resolution, grids, spfun)
# generates comparisonspectra utilizing vimsspec function as well as generates an array of viewing geometries to be utilized in the fitter function.
# 
# parameters:
# 
# cubes = a tuple of cubes that are being utilized to fit.
# coords = a coordinate grid.
# resolution = resolution of the coordinate grid. size of the coordinate grid.
# grids = a tuple of grids that are being utilized. should be one grid for each cube in the same index as its respective cube.
# spfun = interpolator function generated by interp function.
# 
# returns:
# 
# compspec = a n-d array of comparison spectra from the vims database generated by the vimsspec function.
# vgeo = a n-d array of viewing geometries for respective spectra.

# In[ ]:


def geospec(cubes, coords, resolution, grids, spfun):
    comparisonspectra = []
    viewinggeo = []
    for l in range(len(cubes)):
        for a in range(resolution):
            V1643411254, V1643411254spectra = vimsspec(cubes[l], coords[a][1], coords[a][0], grids[l])
            viewinggeo.append(V1643411254)
            comparisonspectra.append(V1643411254spectra)
    comparisonspectra1 = np.array_split(comparisonspectra, len(cubes))
    viewinggeo1 = np.array_split(viewinggeo, len(cubes))
    return comparisonspectra1, viewinggeo1


# fitter documentation:
# specfromvims, generatedspec, params = fitter(geos, compspec, resolution, spfun, channels, la=0.01, ua=0.5, lh=0.1, uh=2, lm=0.4, um=1.4, ia=0.01, ih=1.0, im=1.0)
# fits any number of given spectra utilizing given interpolator function to produce generated spectra, surface albedo, haze scaling factor, and methane scaling factor.
# 
# parameters:
# 
# geos = vgeo returned from geospec function.
# compspec = compspec returned from geospec function.
# resolution = resolution of prior coordinate grid.
# spfun = interpolator function generated by interp function.
# channels = a boolean array of wavelength channels to fit for. 
# la, ua, ia (optional) = lower albedo bound and upper albedo bound. default to 0.01 and 0.5. initial albedo guess. default at 0.01.
# lh, uh, ih (optional) = lower haze bound and upper haze bound. default to 0.1 and 2. initial haze guess. default at 1.0.
# lm, um, im (optional) = lower methane bound and upper methane bound. default to 0.4 and 1.4. initial methane guess. default at 1.0.
# 
# returns:
# 
# specfromvims = flattened array of spectra pulled from the vims database and weighted by euclidean distance.
# generatedspec = spectra generated to fit given best fit parameters.
# params = array with best fit albedo, haze, and methane values for given spectra.

# In[ ]:


def fitter(geos, compspec, resolution, spfun, channels, wv, la=0.01, ua=0.5, lh=0.1, uh=2, lm=0.4, um=1.4, ia=0.01, ih=1.0, im=1.0):
    specfromvims = []
    generatedspec = []
    params = []
    for i in range(resolution):
        spectra=[]
        unflatspec = []
        for a in range(len(compspec)):
            spectra.append(compspec[a][i])
        for b in spectra:
            unflatspec.append(b[channels])
        flatspec = np.ravel(unflatspec)
        def specfit(wave, albedo, haze, methane):
            total = []
            for a in range(len(geos)):
                deltaaz = deltaAZ(geos[a][i]['emission'],geos[a][i]['incidence'],geos[a][i]['phase'])
                valsar = spfun((geos[a][i]['incidence'], geos[a][i]['emission'], deltaaz, albedo, haze, methane, wave))
                total.append(valsar)
            total2 = np.ravel(total)
            return total2
        check3, checkcov = optimize.curve_fit(specfit, wv[channels], flatspec, p0=(ia, ih, im), bounds=((la, lh, lm), (ua,uh,um)))
        genspec = specfit(wv[channels], *check3)
        specfromvims.append(flatspec)
        generatedspec.append(genspec)
        params.append(check3)
    return specfromvims, generatedspec, params

