#!/usr/bin/env python
# Wenchang Yang (wenchang@princeton.edu)
# Thu Sep 19 15:57:34 EDT 2019
from datetime import datetime
import os.path, sys, os
import xarray as xr, numpy as np, pandas as pd
import matplotlib.pyplot as plt
from scipy.optimize import curve_fit

def func(x, x1, x2, x3, y1, y2, y3, k0, k4):
    condlist = [x<=x1, (x>x1)&(x<=x2), (x>x2)&(x<=x3), x>x3]
    funclist = [lambda x: y1 + k0*(x-x1),
        lambda x: y1 + (y2-y1)/(x2-x1)*(x-x1),
        lambda x: y2 + (y3-y2)/(x3-x2)*(x-x2),
        lambda x: y3 + k4*(x-x3) ]

    return np.piecewise(x, condlist, funclist)

def func_k(x, x1, x2, x3, k1, k2, k3, k4, y1):
    condlist = [x<=x1, (x>x1)&(x<=x2), (x>x2)&(x<=x3), x>x3]
    funclist = [lambda x: y1 + k1*(x-x1),
        lambda x: y1 + k2*(x-x1),
        lambda x: y1 + k2*(x2-x1) + k3*(x-x2),
        lambda x: y1 + k2*(x2-x1) + k3*(x3-x2) + k4*(x-x3) ]

    return np.piecewise(x, condlist, funclist)

def func_plinear(x, x1, x2, k1, k2, k3, y1):
    condlist = [x<=x1, (x>x1)&(x<=x2), x>x2]
    funclist = [lambda x: y1 + k1*(x-x1),
        lambda x: y1 + k2*(x-x1),
        lambda x: y1 + k2*(x2-x1) + k3*(x-x2)]

    return np.piecewise(x, condlist, funclist)

def func_linear(x, k, b):
    return k*x + b

def func_exp(x, T0, a, tao):
    return T0 + a*np.exp(-(x-100)/tao) 
 
if __name__ == '__main__':
    tformat = '%Y-%m-%d %H:%M:%S'
    t0 = datetime.now()
    print('[start]:', t0.strftime(tformat))

    ifile = 'data/m6p0sol_CTL1860_tigercpu_intelmpi_18_576PE.t_surf.glbMean.0101-0597.nc'
    da = xr.open_dataarray(ifile).groupby('time.year').mean('time')
    xdata = da.year.values.astype('float')
    ydata = da.values
    p0 = [105, 400, 550,  
        da.sel(year=105).item(), 
        da.sel(year=400).item(), 
        da.sel(year=550).item(), 
        -0.5, -0.5]
    pk0 = [110, 400, 550, -0.5, -0.5, -0.5, -0.5, da.sel(year=110).item()]

    print('piecewise linear')
    #p, e = curve_fit(func, xdata, ydata, p0=p0) 
    #print(p)
    pk, e = curve_fit(func_k, xdata, ydata, p0=pk0) 
    pnames = ('x1', 'x2', 'x3', 'k1', 'k2', 'k3', 'k4', 'y1') 
    print(dict(zip(pnames, pk)))
    print()

    plt.figure(figsize=(8,4))
    da.plot(color='k', label='raw')
    #plt.plot(xdata, func(xdata, *p), label='piecewise1')
    #plt.plot(xdata, func_k(xdata, *pk), label='piecewise linear')

    print('linear 0101-0110')
    years = slice(101,110)
    xdata = da.sel(year=years).year.values.astype('float')
    ydata = da.sel(year=years).values
    p, e = curve_fit(func_linear, xdata, ydata)
    print(p)
    print()
    #plt.plot(xdata, func_linear(xdata, *p), label='linear over 0101-0110')
    
    print('exp 0101-0120')
    years = slice(101,120)
    xdata = da.sel(year=years).year.values.astype('float')
    ydata = da.sel(year=years).values
    p, e = curve_fit(func_exp, xdata, ydata, p0=[282, 5, 3])
    print(p)
    print()
    plt.plot(xdata, func_exp(xdata, *p), label='exp over 0101-0120')


    print('piecewise lienar after 0120')
    years = slice(121, None)
    xdata = da.sel(year=years).year.values.astype('float')
    ydata = da.sel(year=years).values
    p, e = curve_fit(func_plinear, xdata, ydata, 
        p0=[400, 550, -0.5, -0.5, -0.5, da.sel(year=400).item()])
    print(p)
    print()
    plt.plot(xdata, func_plinear(xdata, *p), label='piecewise linear after 0120')

    print('linear 0201-0400')
    years = slice(201,400)
    xdata = da.sel(year=years).year.values.astype('float')
    ydata = da.sel(year=years).values
    p, e = curve_fit(func_linear, xdata, ydata)
    print(p)
    print()


    print('linear 0451-0550')
    years = slice(451,550)
    xdata = da.sel(year=years).year.values.astype('float')
    ydata = da.sel(year=years).values
    p, e = curve_fit(func_linear, xdata, ydata)
    print(p)
    print()


    print('linear 0561-0590')
    years = slice(561,590)
    xdata = da.sel(year=years).year.values.astype('float')
    ydata = da.sel(year=years).values
    p, e = curve_fit(func_linear, xdata, ydata)
    print(p)
    print()

    plt.legend()
    plt.tight_layout()
    #plt.savefig('fig_curve_fit.png')
    plt.show()

    t1 = datetime.now()
    print('[end]:', t1.strftime(tformat))
    print('[total time]:', f'{(t1-t0).seconds:,} seconds')