import pandas as pd
import os
from netCDF4 import Dataset
import numpy as np
from datetime import timedelta,datetime
import folium
import matplotlib
import matplotlib.pyplot as plt

def find_nearest(lat,lon):
    lat_lon_idx = np.where((lats-lat)**2+(lons-lon)**2==np.min((lats-lat)**2+(lons-lon)**2))
    lat_idx = lat_lon_idx[0][0]
    lon_idx = lat_lon_idx[1][0]
    return lat_idx,lon_idx
def calculate_new_position(lat, lon, u, v):
    # Calculate displacement in degrees latitude and longitude
    # Assuming 1 degree latitude ~ 111 km
    # and 1 degree longitude varies based on latitude (using an average value here for simplicity)
    lat_displacement = (-1)*v * 3600 / 111000  # v in m/s, displacement in degrees
    lon_displacement = (-1)*u * 3600 / (111000 * np.cos(np.radians(lat)))
    return lat + lat_displacement, lon + lon_displacement
# Define helper functions
# def get_wrf_data(dt, lat, lon):
    # Adjust the path and filename as needed for your WRF file structure
#     filepath = f"/expanse/lustre/scratch/xhu2/temp_project/Run/CO2_and_otherGHG/WRFV4.2.2/CONUS/wrfchem4.2.2_R2_CONUS_Hu2021JGR_CAMSCO2only.{dt.strftime('%Y')}010100"
#     filename = f"wrfout_d01_{dt.strftime('%Y-%m-%d_%H:00:00')}"

#     lat_idx, lon_idx = find_nearest(lat, lon)
    # print(lat_idx,lon_idx)
    # Here, implement the logic to open the WRF file and extract U10, V10 for the given lat, lon
#     wrfout_file = os.path.join(filepath,filename)
#     fid = Dataset(wrfout_file)
    # print("U10 shape")
    # print(np.shape(fid['U10'][:]))
#     u10 = fid['U10'][0,lat_idx,lon_idx]
#     v10 = fid['V10'][0,lat_idx,lon_idx]
#     fid.close()
#     return u10, v10
#def normalize_index(value):
#    rounded_value = round(value / 50) * 50
#    return int((rounded_value-100)/50)
# Normalize function to get an index for the colors array
def normalize_index(value, min_value=500, max_value=1000, step=25):
    # Round the value to the nearest multiple of the step
    rounded_value = round((value - min_value) / step) * step
    # Calculate the index in the colors array
    if value>max_value:
        index = -1
    else:
    # Calculate the index in the colors array
        index = int((rounded_value - min_value) / step)
        # Ensure the index is within the bounds of the colors array
        index = max(0, min(index, len(colors_month) - 1))
    return index
def round_down_to_nearest_three_hours(dt):
    rounded_hour = dt.hour - (dt.hour % 3)
    return dt.replace(hour=rounded_hour, minute=0, second=0, microsecond=0)

def get_wrf_data(dt, lat, lon):
    dt_rounded = round_down_to_nearest_three_hours(dt)
    if dt.year == 2016:
        filepath = "/expanse/lustre/scratch/xhu2/temp_project/Run/CO2_and_otherGHG/WRFV4.2.2/CONUS/wrfchem4.2.2_R2_CONUS_Sharon_CH4NEI2017_CAMScorrect_Wetland131_agwasteOce.2016010100_outputs"
    elif dt.year == 2017:
        filepath = "/expanse/lustre/projects/uok114/xhu2/FromComet/scratch/Run/CO2_and_otherGHG/WRFV4.1.5/CONUS/wrfchem4.1.5_R2_CONUS_trac16_10mb_Hil3ReShrubRES_addOce_restoreDF_ODIAC_CT2019.2017010100"
    elif dt.year == 2018:
        filepath = "/expanse/lustre/scratch/xhu2/temp_project/Run/CO2_and_otherGHG/WRFV4.2.2/CONUS/wrfchem4.2.2_R2_CONUS_Hu2021JGR_CT2019_CH4NEI2017_CAMS_Wetland131_agwasteOce.2018010100"
    elif dt.year == 2019:
        filepath = "/expanse/lustre/scratch/xhu2/temp_project/Run/CO2_and_otherGHG/WRFV4.2.2/CONUS/wrfchem4.2.2_R2_CONUS_Hu2021JGR_CH4NEI2017_CAMSboth_Wetland131_agwasteOce.2019010100"
    elif dt.year == 2020:
        filepath = "/expanse/lustre/scratch/xhu2/temp_project/Run/CO2_and_otherGHG/WRFV4.2.2/CONUS/wrfchem4.2.2_R2_CONUS_Hu2021JGR_CAMSCO2only.2020010100"
    elif dt.year == 2021:
        filepath = f"/expanse/lustre/scratch/xhu2/temp_project/Run/CO2_and_otherGHG/WRFV4.2.2/CONUS/wrfchem4.2.2_R2_CONUS_Hu2021JGR_CAMSCO2only.2021010100"
    else:
        print("Check the year")
    filename = f"wrfout_d01_{dt_rounded.strftime('%Y-%m-%d_%H:%M:%S')}"
    wrfout_file = os.path.join(filepath, filename)
    if os.path.exists(wrfout_file) and os.path.isfile(wrfout_file):
        try:
            fid = Dataset(wrfout_file)
            lat_idx, lon_idx = find_nearest(lat, lon)
            u = fid['U'][0, 2, lat_idx, lon_idx]
            v = fid['V'][0, 2, lat_idx, lon_idx]
            # u = fid['U10'][0, lat_idx, lon_idx]
            # v = fid['V10'][0, lat_idx, lon_idx]
            fid.close()
            return u, v
        except Exception as e:
            print(f"Failed to open or process NetCDF file: {e}")
            return 0, 0
    else:
        print(f"File does not exist or is not a file: {wrfout_file}")
        return 0, 0

selected_pbl_days = ['2012-9-20','2017-11-23','2018-11-15','2019-11-9','2019-11-10','2019-11-24',\
    '2019-11-26','2021-11-23','2021-11-29','2021-11-30','2012-3-12','2013-3-3','2020-3-25','2022-10-19',\
    '2019-6-25','2022-7-11',\
    '2011-12-24','2011-12-29','2012-2-14','2012-12-2','2012-12-11','2012-12-16','2012-12-21',\
    '2012-12-29','2013-1-2','2013-1-19','2013-1-23','2013-2-5','2013-2-8','2013-2-13','2013-2-17',\
    '2013-12-16','2014-1-7','2014-2-12','2014-2-13','2017-2-7','2017-2-23','2017-12-3','2017-12-18',\
    '2018-2-8','2018-2-14','2018-12-5','2018-12-25','2018-12-30','2019-1-4','2019-1-5','2019-1-14','2019-1-15',\
    '2019-1-20','2019-2-1','2019-12-13','2019-12-22','2019-12-23','2020-1-26','2020-1-27','2020-2-2','2020-12-9',\
    '2020-12-10','2020-12-17','2021-1-2','2021-1-3','2021-1-12','2021-1-13','2021-2-20','2021-2-23',\
    '2021-12-21','2021-12-22','2022-3-1','2022-3-2','2022-3-3','2022-1-3','2022-1-4','2022-1-18',\
    '2022-1-31','2022-2-5','2022-2-6','2022-2-8']
# Read the DataFrame
df = pd.read_csv("../data/combined_df.csv", parse_dates=["Time"])
df['Wind_Speed'] = pd.to_numeric(df['Wind_Speed'], errors='coerce')
filtered_events = df
#filtered_events = df[df['xch4_anomaly'] >-999.].copy()
# filtered_events = filtered_events[filtered_events['event']==1]
# print(filtered_events)
# Convert the index to a datetime index if it's not already
filtered_events.index = pd.to_datetime(filtered_events['Time'])
filtered_events = filtered_events[(filtered_events['pblh_sounding']>0)&(filtered_events['pblh_sounding']<500)]
#filtered_events = filtered_events[(filtered_events['Wind_Speed']>1)&(filtered_events['Wind_Speed']<6)]
filtered_events = filtered_events[filtered_events['Height1_anomaly']>300]
#filtered_events = filtered_events[(filtered_events['RiB']>0.25)&(filtered_events['RiB']<1.0)]
print(np.shape(filtered_events))
#filtered_events = filtered_events[filtered_events.index.month.isin([12,1,2])]
#filtered_events = filtered_events.loc[filtered_events.index.normalize().isin(pd.to_datetime(selected_pbl_days))]
# print('1')
# Step 2: Add 6 hours to each time for the back trajectory calculation
filtered_events['time_for_trajectory'] = filtered_events.index + timedelta(hours=6)
#filtered_events = filtered_events[(filtered_events.index.hour>=12)&(filtered_events.index.hour<=18)]
# Step 3: Filter for the first occurrence of continuous events and one event per day
## Only see afternoon anomalies
#filtered_events = filtered_events[filtered_events['PBLH']>200]
# filtered_events = filtered_events[(filtered_events.index.hour>10)&(filtered_events.index.hour<21)]
filtered_events['date'] = filtered_events.index.date
filtered_events = filtered_events.drop_duplicates(subset='date')
# print('2')
# print(filtered_events)
filtered_events = filtered_events[(filtered_events.index.year>=2016)&(filtered_events.index.year<=2021)]
# filtered_events = filtered_events[filtered_events.index.hour.isin([12,13,14])]
# print('3')
# print(filtered_events)
# Read universal lat and lin for futyre use
wrfout_file = '../data/wrfout_d01_2021-01-01_01:00:00'
# Open the wrfout file
ds = Dataset(wrfout_file)
# Read the XLAT and XLONG variables
lats = ds.variables['XLAT'][0, :, :]  # Assuming time dimension is first and you want the first time step
lons = ds.variables['XLONG'][0, :, :]  # Same assumption as above
# Close the dataset after extraction
ds.close()
# print('4')

# Initialize map
map = folium.Map(location=[36.6058, -97.4888], zoom_start=6)

# Define colors for different events
colors = plt.cm.viridis(np.linspace(0, 1, len(filtered_events)))
#colors_month = plt.cm.jet(np.linspace(0, 1, 35))
# Generate an array of colors from the jet colormap
num_colors = int((300 - 100) / 25) + 1  # One color for each 50 unit interval from 0 to 1700
colors_month = [plt.cm.jet(i / (num_colors - 1)) for i in range(num_colors)]

# Loop through each event
for index, (color, (_, event)) in enumerate(zip(colors, filtered_events.iterrows())):
    lat, lon = 36.6058, -97.4888  # Starting coordinates
    dt = event['time_for_trajectory']
    value = event['Height1_anomaly']
    print(dt)
    
    trajectory = [(lat, lon)]
    for i in range(6):  # Go back 24 hours
        u, v = get_wrf_data(dt - timedelta(hours=i), lat, lon)
        lat, lon = calculate_new_position(lat, lon, u, v)
        trajectory.append((lat, lon))
    
    # Plot trajectory on map
    # plt.cm.jet(normalized_value)
    # matplotlib.colors.to_hex(colors_month[dt.month-1])
    folium.PolyLine(trajectory, color=matplotlib.colors.to_hex(colors_month[normalize_index(value)]), weight=2.5, opacity=1).add_to(map)

# Finalize the map
folium.Marker([36.6058, -97.4888], popup='Starting Point', icon=folium.Icon(color='red')).add_to(map)
folium.Marker([36.641160, -97.536859], popup='Cattle Farm', icon=folium.Icon(color='blue')).add_to(map)
map.save('back_trajectory_map_0610_1.html')