import pandas as pd
import numpy as np
import seaborn as sns
import matplotlib.pyplot as plt
import matplotlib.dates as mdates
import scipy.stats as stats

from datetime import datetime, timedelta
from scipy.optimize import curve_fit

%matplotlib inline

# Seaborn defaults
sns.set_theme()
sns.set(style = 'white', font_scale = 1.5, rc = {'figure.figsize':(16, 9)})

bee = pd.read_csv('../../Google Drive/job_tech_learning/data/Bumblebeesandfl/MT_2023_Bee_Plant_Data.csv')
bee.head()

bee.shape

rename_cols = {
    "Site ID": 'id', 
    "Date": 'date_str', 
    "Bee Survey Length (min)": "survey_length", 
    "Number of Bumble Bees Captured": 'n_captured', 
    "Number of Plant Species in Flower": 'n_flowering_plants', 
    "List of Plant Species in Flower": 'sp_str'
}

bee = (
    bee.
    rename(columns = rename_cols).
    assign(
        date = lambda x: pd.to_datetime(x.date_str, format = "%m/%d/%Y")
    )
)

bee

# Split spp string into list of strings

bee = (
    bee.
    assign(
        sp_list = lambda x: [i.split(", ") for i in x.sp_str]
    )
)

bee.sort_values(['id', 'date'])

# not quite right; the length of every sp_list shouldn't be 65...

[len(i) for i in bee.sp_list][:5]

# Ah. The sp_list is just the list of the flowering plants and will equal the n_flowering_plants col.
# it's not a list of all plants in the sampling area. Bummer.

fig, ax = plt.subplots()


g = sns.scatterplot(data = bee, x = 'date', y = 'n_flowering_plants')

# g.axvline(35, color = 'black', ls = "--")
g.set(
    ylabel = 'N flowering plants in sample area',
    xlabel = 'Date'
    )#xlim = (-25, 730))

# Define the model
def von_b(x, Linf, K, t0): 
    return Linf * (1 - np.exp(-K*(x-t0)))

# optimized values, covariance of optimized values
params, cov = curve_fit(
    f = von_b, 
    xdata = mdates.date2num(bee.date), 
    ydata = bee.n_flowering_plants, 
    p0 = [19, 0, 0]
)

print(params, "\n\n", cov)

xdata = mdates.date2num(bee.date)
ydata = bee.n_flowering_plants

fig, ax = plt.subplots()

# g = sns.scatterplot(xdata, von_b(xdata, *res), 'r-', label = 'fit')
sns.scatterplot(x = bee.date, y = ydata, label = 'data')
sns.lineplot(x = xdata, y = von_b(xdata, *params), c = 'red', label = 'fit')
# g # plt.show()

# Get CI, t-stat, and p-value for parameters from curve_fit() 
# from: https://stats.stackexchange.com/questions/362520/how-to-know-if-a-parameter-is-statistically-significant-in-a-curve-fit-estimat

import scipy.odr
import scipy.stats

def f_wrapper_for_odr(params, x): # parameter order for odr
    return von_b(x, *params)

model = scipy.odr.odrpack.Model(f_wrapper_for_odr)

data = scipy.odr.odrpack.Data(xdata, ydata)
myodr = scipy.odr.odrpack.ODR(data, model, beta0 = params, maxit = 0)

myodr.set_job(fit_type = 2)

parameterStatistics = myodr.run()
df_e = len(xdata)-len(params)
cov_beta = parameterStatistics.cov_beta
sd_beta = parameterStatistics.sd_beta * parameterStatistics.sd_beta

t_df = scipy.stats.t.ppf(0.975, df_e) #identify dof for 95% CI

ci = [] #Empty list for storage
for i in range(len(params)):
               ci.append([params[i] - t_df * parameterStatistics.sd_beta[i], params[i] + t_df * parameterStatistics.sd_beta])
        
tstat_beta = params / parameterStatistics.sd_beta
pstat_beta = (1.0 - scipy.stats.t.cdf(np.abs(tstat_beta), df_e)) * 2.0

model.

ci

print(
    params, "\n", 
    tstat_beta, "\n", 
    pstat_beta
)

for i in range(len(params)):
    print('parameter:', params[i])
    print('   conf interval:', ci[i][0], ci[i][1])
    print('   t-statistic:', tstat_beta[i])
    print('   p-value:', pstat_beta[i])
    print()

import matplotlib.dates as mdates 

mdates.date2num(bee.date)

# import numpy as np
# from scipy.optimize import curve_fit
# import matplotlib.pyplot as plt

# def f(x, start, end):
#     res = np.empty_like(x)
#     res[x < start] =-1
#     res[x > end] = 1
#     linear = np.all([[start <= x], [x <= end]], axis=0)[0]
#     res[linear] = np.linspace(-1., 1., num=np.sum(linear))
#     return res

# if __name__ == '__main__':

#     xdata = np.linspace(0., 1000., 1000)
#     ydata = -np.ones(1000)
#     ydata[500:1000] = 1.
#     ydata = ydata + np.random.normal(0., 0.25, len(ydata))

#     popt, pcov = curve_fit(f, xdata, ydata, p0=[495., 505.])
#     print(popt, pcov)
#     plt.figure()
#     plt.plot(xdata, f(xdata, *popt), 'r-', label='fit')
#     plt.plot(xdata, ydata, 'b-', label='data')
#     plt.show()

np.empty_like(bee.date)

bee = pd.read_csv('../../Google Drive/job_tech_learning/data/Bumblebeesandfl/MT_2023_Plant_Species.csv')
bee.head()