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'


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

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

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

# Convert temperature to numeric
fin['temp'] = pd.to_numeric(fin['temp'])

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)

temp_mean = temp.mean(dim='time', skipna=True).item()
b = temp_mean
print("Mean temperature, b =", b)

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

print("a =", a)

y = a * np.cos(w * (t - phi)) + b


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)

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

#ax.plot(date, temp, 'o', markersize=2, color='C1')
ax.plot(t, temp, 'o', markersize=2, color='C1')
ax.plot(t, y, '-', color='k', linewidth=2, label='Estimation')

title_name = "Seasonal Cycle Fit"

ax.set_xlabel('Time [days]')
ax.set_ylabel('Air 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()