import os
import numpy as np
import pandas as pd
import xarray as xr
import matplotlib.pyplot as plt
import matplotlib.dates as mdates

# ----------------------------
# Settings
# ----------------------------
#file_csv = 'greenfarm_temp.csv'
#title_name = 'Greenville Farm Weather Station'
#file_csv = 'Venice_temp.csv'
#title_name = 'Venice Weather Station'
file_csv = 'DugwayUT.csv'
title_name = 'Dugway, UT'


# ----------------------------
# Read CSV file
# ----------------------------
print(' ----- Read CSV file ------')

fin = pd.read_csv(
    file_csv,
    na_values=['M', 'T'],
    encoding='ISO-8859-1',
    delimiter=',',
    header=71,
    #names=['date', 'temp']
)

# Convert timestamp to datetime
fin['date'] = pd.to_datetime(fin['Date'])

# Convert temperature to numeric
fin['temp'] = pd.to_numeric(fin['Soil Temperature Observed -4in (degF) Mean of Hourly Values'])

print(fin)
print(' ----------------------------')

# Show missing-value count
print(' ----- Missing value check ------')
print(fin.isna().sum())
print(' ----------------------------')

# ----------------------------
# Time and temperature setting
# ----------------------------
print(' ----- Time setting ------')
date = fin['date']
t = xr.DataArray(range(len(fin)), dims=('time',), name='time', coords={'date': ('time', date)})
print(date)
print(' ----------------------------')

print(' ----- Temperature setting ------')
#temp = xr.DataArray(fin['temp'], dims=('date',), coords={'date': date})
temp = xr.DataArray(fin['temp'], dims=('time',), coords={'time': t, 'date': ('time', date)})
print(temp)
print(' ----------------------------')


w = 2.0 * np.pi / 365.0
phi = 15.0

#a = -20.0
#b = 50.0

#t_mod = t % 365
#
#temp_summer = temp.where((t_mod >= 190) & (t_mod <= 200), drop=True)
#temp_winter = temp.where((t_mod >= 10) & (t_mod <= 20), drop=True)
#
#T_summer = temp_summer.mean(dim='time', skipna=True).item()
#T_winter = temp_winter.mean(dim='time', skipna=True).item()
#
#print("Mid-summer temperature =", T_summer)
#print("Mid-winter temperature =", T_winter)
#
#b = 0.5 * (T_summer + T_winter)
#a = -0.5 * (T_summer - T_winter)

t_mean = t.mean().item()
temp_mean = temp.mean(dim='time', skipna=True).item()


f = np.cos(w * (t - phi))
#a = ((temp - b) * f).mean(dim='time', skipna=True).item() / (f * f).mean(dim='time', skipna=True).item()


#c = 8.0 / (5.0 * 365.0)
#f1 = t
#c = ((temp - temp_mean) * f1).mean(dim='time', skipna=True).item() / (f1 * f1).mean(dim='time', skipna=True).item()
#print("c =", c)
#
#b = temp_mean - c * t_mean
#print("Mean temperature, b =", b)

coef = np.polyfit(t.values, temp.values, 1)
c = coef[0]
b = coef[1]

a = ((temp - b - c * t) * f).mean(dim='time', skipna=True).item() / (f * f).mean(dim='time', skipna=True).item()
print("a =", a)

y1 = a * np.cos(w * (t - phi)) + b
y2 = a * np.cos(w * (t - phi)) + b + c * t
y_trend = c * t


#se = (temp - y) * (temp - y)
#mse = se.mean(dim='time', skipna=True).item()
#rmse = np.sqrt(mse)
# OR one line version
# rmse = np.sqrt(((temp - y)**2).mean(dim='time', skipna=True).item())
#print("RMSE =", rmse)

temp_anm1 = temp - y1
temp_anm2 = temp - y2


# ----------------------------
# Plot
# ----------------------------
fig, ax = plt.subplots()

#ax.plot(date, temp, 'o', markersize=2, color='C1')
#ax.plot(t, temp, 'o', markersize=2, color='C1', label='Observation')
#ax.plot(t, temp_rm, '-', linewidth=2, color='C0', label='Running mean')
#ax.plot(t, y, '-', color='k', linewidth=2, label='Estimation')
#ax.plot(t, temp_anm, 'o', markersize=2, color='C1', label='Anomaly')
#ax.plot(t, temp_anm_rm, '-', linewidth=2, color='C0', label='Running mean')

ax.plot(t, temp_anm1, 'o', markersize=2, color='pink', label='No seasonal cycle')
#ax.plot(t, y_trend, '-', color='k', linewidth=2, label='Linear trend')
ax.plot(t, temp_anm2, 'o', markersize=2, color='lightblue', label='No seasonal cycle & detrend')

temp_anm1_rm = temp_anm1.rolling(time=365, center=True).mean()
temp_anm2_rm = temp_anm2.rolling(time=365, center=True).mean()

#ax.plot(t, y_trend, '-', color='k', linewidth=2, label='Linear trend')
ax.plot(t, temp_anm1_rm, '-', linewidth=2, color='red', label='Running mean 1')
ax.plot(t, temp_anm2_rm, '-', linewidth=2, color='darkblue', label='Running mean 2')

title_name = "Seasonal Cycle Fit"

ax.set_xlabel('Time [days]')
ax.set_ylabel('Soil Temperature [°F]')
#ax.set_title(f'{title_name}\nRMSE = {rmse:.2f}°F')

ax.set_title(title_name)

#ax.text( 0.98, 0.95, f'RMSE = {rmse:.2f}°F', transform=ax.transAxes, ha='right', va='top')

ax.legend()

ax.xaxis.set_major_locator(mdates.YearLocator())
ax.xaxis.set_major_formatter(mdates.DateFormatter('%Y-%m-%d'))
ax.xaxis.set_minor_locator(mdates.MonthLocator())

#ax.grid(True)
ax.grid(which='major', linewidth=1)
ax.grid(which='minor', linestyle='--', alpha=0.4)

fig.autofmt_xdate()


# ----------------------------
# Save figure using script name
# ----------------------------
script_name = os.path.splitext(os.path.basename(__file__))[0]
outfile = script_name + '.png'

plt.savefig(outfile)
print('Figure saved as:', outfile)

plt.show()