import gradio as gr
import pandas as pd
import numpy as np
import plotly.graph_objects as go
import plotly.express as px
import tropycal.tracks as tracks
import pickle
import requests
import os
import argparse
from datetime import datetime
import statsmodels.api as sm
import shutil
import tempfile
import csv
from collections import defaultdict
import filecmp

# Command-line argument parsing
parser = argparse.ArgumentParser(description='Typhoon Analysis Dashboard')
parser.add_argument('--data_path', type=str, default=os.getcwd(), help='Path to the data directory')
args = parser.parse_args()
DATA_PATH = args.data_path

ONI_DATA_PATH = os.path.join(DATA_PATH, 'oni_data.csv')
TYPHOON_DATA_PATH = os.path.join(DATA_PATH, 'processed_typhoon_data.csv')
LOCAL_iBtrace_PATH = os.path.join(DATA_PATH, 'ibtracs.WP.list.v04r01.csv')
iBtrace_uri = 'https://www.ncei.noaa.gov/data/international-best-track-archive-for-climate-stewardship-ibtracs/v04r01/access/csv/ibtracs.WP.list.v04r01.csv'
CACHE_FILE = 'ibtracs_cache.pkl'
CACHE_EXPIRY_DAYS = 1

# Color map for typhoon categories
color_map = {
    'C5 Super Typhoon': 'rgb(255, 0, 0)',
    'C4 Very Strong Typhoon': 'rgb(255, 63, 0)',
    'C3 Strong Typhoon': 'rgb(255, 127, 0)',
    'C2 Typhoon': 'rgb(255, 191, 0)',
    'C1 Typhoon': 'rgb(255, 255, 0)',
    'Tropical Storm': 'rgb(0, 255, 255)',
    'Tropical Depression': 'rgb(173, 216, 230)'
}

# Classification standards
atlantic_standard = {
    'C5 Super Typhoon': {'wind_speed': 137, 'color': 'rgb(255, 0, 0)'},
    'C4 Very Strong Typhoon': {'wind_speed': 113, 'color': 'rgb(255, 63, 0)'},
    'C3 Strong Typhoon': {'wind_speed': 96, 'color': 'rgb(255, 127, 0)'},
    'C2 Typhoon': {'wind_speed': 83, 'color': 'rgb(255, 191, 0)'},
    'C1 Typhoon': {'wind_speed': 64, 'color': 'rgb(255, 255, 0)'},
    'Tropical Storm': {'wind_speed': 34, 'color': 'rgb(0, 255, 255)'},
    'Tropical Depression': {'wind_speed': 0, 'color': 'rgb(173, 216, 230)'}
}

taiwan_standard = {
    'Strong Typhoon': {'wind_speed': 51.0, 'color': 'rgb(255, 0, 0)'},
    'Medium Typhoon': {'wind_speed': 33.7, 'color': 'rgb(255, 127, 0)'},
    'Mild Typhoon': {'wind_speed': 17.2, 'color': 'rgb(255, 255, 0)'},
    'Tropical Depression': {'wind_speed': 0, 'color': 'rgb(173, 216, 230)'}
}

# Data loading and preprocessing functions
def download_oni_file(url, filename):
    response = requests.get(url)
    response.raise_for_status()
    with open(filename, 'wb') as f:
        f.write(response.content)
    return True

def convert_oni_ascii_to_csv(input_file, output_file):
    data = defaultdict(lambda: [''] * 12)
    season_to_month = {'DJF': 12, 'JFM': 1, 'FMA': 2, 'MAM': 3, 'AMJ': 4, 'MJJ': 5,
                       'JJA': 6, 'JAS': 7, 'ASO': 8, 'SON': 9, 'OND': 10, 'NDJ': 11}
    with open(input_file, 'r') as f:
        lines = f.readlines()[1:]
        for line in lines:
            parts = line.split()
            if len(parts) >= 4:
                season, year, anom = parts[0], parts[1], parts[-1]
                if season in season_to_month:
                    month = season_to_month[season]
                    if season == 'DJF':
                        year = str(int(year) - 1)
                    data[year][month-1] = anom
    with open(output_file, 'w', newline='') as f:
        writer = csv.writer(f)
        writer.writerow(['Year', 'Jan', 'Feb', 'Mar', 'Apr', 'May', 'Jun', 'Jul', 'Aug', 'Sep', 'Oct', 'Nov', 'Dec'])
        for year in sorted(data.keys()):
            writer.writerow([year] + data[year])

def update_oni_data():
    url = "https://www.cpc.ncep.noaa.gov/data/indices/oni.ascii.txt"
    temp_file = os.path.join(DATA_PATH, "temp_oni.ascii.txt")
    input_file = os.path.join(DATA_PATH, "oni.ascii.txt")
    output_file = ONI_DATA_PATH
    if download_oni_file(url, temp_file):
        if not os.path.exists(input_file) or not filecmp.cmp(temp_file, input_file):
            os.replace(temp_file, input_file)
            convert_oni_ascii_to_csv(input_file, output_file)
        else:
            os.remove(temp_file)

def load_ibtracs_data():
    if os.path.exists(CACHE_FILE) and (datetime.now() - datetime.fromtimestamp(os.path.getmtime(CACHE_FILE))).days < CACHE_EXPIRY_DAYS:
        with open(CACHE_FILE, 'rb') as f:
            return pickle.load(f)
    if os.path.exists(LOCAL_iBtrace_PATH):
        ibtracs = tracks.TrackDataset(basin='west_pacific', source='ibtracs', ibtracs_url=LOCAL_iBtrace_PATH)
    else:
        response = requests.get(iBtrace_uri)
        response.raise_for_status()
        with tempfile.NamedTemporaryFile(mode='w', delete=False, suffix='.csv') as temp_file:
            temp_file.write(response.text)
            shutil.move(temp_file.name, LOCAL_iBtrace_PATH)
        ibtracs = tracks.TrackDataset(basin='west_pacific', source='ibtracs', ibtracs_url=LOCAL_iBtrace_PATH)
    with open(CACHE_FILE, 'wb') as f:
        pickle.dump(ibtracs, f)
    return ibtracs

def convert_typhoondata(input_file, output_file):
    with open(input_file, 'r') as infile:
        next(infile); next(infile)  # Skip header lines
        reader = csv.reader(infile)
        sid_data = defaultdict(list)
        for row in reader:
            if row:
                sid = row[0]
                sid_data[sid].append((row, row[6]))
    with open(output_file, 'w', newline='') as outfile:
        fieldnames = ['SID', 'ISO_TIME', 'LAT', 'LON', 'SEASON', 'NAME', 'WMO_WIND', 'WMO_PRES', 'USA_WIND', 'USA_PRES', 'START_DATE', 'END_DATE']
        writer = csv.DictWriter(outfile, fieldnames=fieldnames)
        writer.writeheader()
        for sid, data in sid_data.items():
            start_date = min(data, key=lambda x: x[1])[1]
            end_date = max(data, key=lambda x: x[1])[1]
            for row, iso_time in data:
                writer.writerow({
                    'SID': row[0], 'ISO_TIME': iso_time, 'LAT': row[8], 'LON': row[9], 'SEASON': row[1], 'NAME': row[5],
                    'WMO_WIND': row[10].strip() or ' ', 'WMO_PRES': row[11].strip() or ' ',
                    'USA_WIND': row[23].strip() or ' ', 'USA_PRES': row[24].strip() or ' ',
                    'START_DATE': start_date, 'END_DATE': end_date
                })

def load_data(oni_path, typhoon_path):
    oni_data = pd.read_csv(oni_path)
    typhoon_data = pd.read_csv(typhoon_path, low_memory=False)
    typhoon_data['ISO_TIME'] = pd.to_datetime(typhoon_data['ISO_TIME'], errors='coerce')
    typhoon_data = typhoon_data.dropna(subset=['ISO_TIME'])
    return oni_data, typhoon_data

def process_oni_data(oni_data):
    oni_long = oni_data.melt(id_vars=['Year'], var_name='Month', value_name='ONI')
    month_map = {'Jan': '01', 'Feb': '02', 'Mar': '03', 'Apr': '04', 'May': '05', 'Jun': '06',
                 'Jul': '07', 'Aug': '08', 'Sep': '09', 'Oct': '10', 'Nov': '11', 'Dec': '12'}
    oni_long['Month'] = oni_long['Month'].map(month_map)
    oni_long['Date'] = pd.to_datetime(oni_long['Year'].astype(str) + '-' + oni_long['Month'] + '-01')
    oni_long['ONI'] = pd.to_numeric(oni_long['ONI'], errors='coerce')
    return oni_long

def process_typhoon_data(typhoon_data):
    typhoon_data['ISO_TIME'] = pd.to_datetime(typhoon_data['ISO_TIME'], errors='coerce')
    typhoon_data['USA_WIND'] = pd.to_numeric(typhoon_data['USA_WIND'], errors='coerce')
    typhoon_data['USA_PRES'] = pd.to_numeric(typhoon_data['USA_PRES'], errors='coerce')
    typhoon_data['LON'] = pd.to_numeric(typhoon_data['LON'], errors='coerce')
    typhoon_max = typhoon_data.groupby('SID').agg({
        'USA_WIND': 'max', 'USA_PRES': 'min', 'ISO_TIME': 'first', 'SEASON': 'first', 'NAME': 'first',
        'LAT': 'first', 'LON': 'first'
    }).reset_index()
    typhoon_max['Month'] = typhoon_max['ISO_TIME'].dt.strftime('%m')
    typhoon_max['Year'] = typhoon_max['ISO_TIME'].dt.year
    typhoon_max['Category'] = typhoon_max['USA_WIND'].apply(categorize_typhoon)
    return typhoon_max

def merge_data(oni_long, typhoon_max):
    return pd.merge(typhoon_max, oni_long, on=['Year', 'Month'])

def categorize_typhoon(wind_speed):
    wind_speed_kt = wind_speed  # Assuming input is already in knots
    if wind_speed_kt >= 137:
        return 'C5 Super Typhoon'
    elif wind_speed_kt >= 113:
        return 'C4 Very Strong Typhoon'
    elif wind_speed_kt >= 96:
        return 'C3 Strong Typhoon'
    elif wind_speed_kt >= 83:
        return 'C2 Typhoon'
    elif wind_speed_kt >= 64:
        return 'C1 Typhoon'
    elif wind_speed_kt >= 34:
        return 'Tropical Storm'
    else:
        return 'Tropical Depression'

def classify_enso_phases(oni_value):
    if isinstance(oni_value, pd.Series):
        oni_value = oni_value.iloc[0]
    if oni_value >= 0.5:
        return 'El Nino'
    elif oni_value <= -0.5:
        return 'La Nina'
    else:
        return 'Neutral'

# Load data globally
update_oni_data()
ibtracs = load_ibtracs_data()
convert_typhoondata(LOCAL_iBtrace_PATH, TYPHOON_DATA_PATH)
oni_data, typhoon_data = load_data(ONI_DATA_PATH, TYPHOON_DATA_PATH)
oni_long = process_oni_data(oni_data)
typhoon_max = process_typhoon_data(typhoon_data)
merged_data = merge_data(oni_long, typhoon_max)

# Main analysis functions
def generate_typhoon_tracks(filtered_data, typhoon_search):
    fig = go.Figure()
    for sid in filtered_data['SID'].unique():
        storm_data = filtered_data[filtered_data['SID'] == sid]
        color = {'El Nino': 'red', 'La Nina': 'blue', 'Neutral': 'green'}[storm_data['ENSO_Phase'].iloc[0]]
        fig.add_trace(go.Scattergeo(
            lon=storm_data['LON'], lat=storm_data['LAT'], mode='lines',
            name=storm_data['NAME'].iloc[0], line=dict(width=2, color=color)
        ))
    if typhoon_search:
        mask = filtered_data['NAME'].str.contains(typhoon_search, case=False, na=False)
        if mask.any():
            storm_data = filtered_data[mask]
            fig.add_trace(go.Scattergeo(
                lon=storm_data['LON'], lat=storm_data['LAT'], mode='lines',
                name=f'Matched: {typhoon_search}', line=dict(width=5, color='yellow')
            ))
    fig.update_layout(title='Typhoon Tracks', geo=dict(projection_type='natural earth', showland=True))
    return fig

def generate_wind_oni_scatter(filtered_data, typhoon_search):
    fig = px.scatter(filtered_data, x='ONI', y='USA_WIND', color='Category', hover_data=['NAME', 'Year', 'Category'],
                     title='Wind Speed vs ONI', labels={'ONI': 'ONI Value', 'USA_WIND': 'Max Wind Speed (knots)'},
                     color_discrete_map=color_map)
    if typhoon_search:
        mask = filtered_data['NAME'].str.contains(typhoon_search, case=False, na=False)
        if mask.any():
            fig.add_trace(go.Scatter(
                x=filtered_data.loc[mask, 'ONI'], y=filtered_data.loc[mask, 'USA_WIND'],
                mode='markers', marker=dict(size=10, color='red', symbol='star'),
                name=f'Matched: {typhoon_search}',
                text=filtered_data.loc[mask, 'NAME'] + ' (' + filtered_data.loc[mask, 'Year'].astype(str) + ')'
            ))
    return fig

def generate_pressure_oni_scatter(filtered_data, typhoon_search):
    fig = px.scatter(filtered_data, x='ONI', y='USA_PRES', color='Category', hover_data=['NAME', 'Year', 'Category'],
                     title='Pressure vs ONI', labels={'ONI': 'ONI Value', 'USA_PRES': 'Min Pressure (hPa)'},
                     color_discrete_map=color_map)
    if typhoon_search:
        mask = filtered_data['NAME'].str.contains(typhoon_search, case=False, na=False)
        if mask.any():
            fig.add_trace(go.Scatter(
                x=filtered_data.loc[mask, 'ONI'], y=filtered_data.loc[mask, 'USA_PRES'],
                mode='markers', marker=dict(size=10, color='red', symbol='star'),
                name=f'Matched: {typhoon_search}',
                text=filtered_data.loc[mask, 'NAME'] + ' (' + filtered_data.loc[mask, 'Year'].astype(str) + ')'
            ))
    return fig

def generate_regression_analysis(filtered_data):
    fig = px.scatter(filtered_data, x='LON', y='ONI', hover_data=['NAME'],
                     title='Typhoon Generation Longitude vs ONI (All Years)')
    if len(filtered_data) > 1:
        X = np.array(filtered_data['LON']).reshape(-1, 1)
        y = filtered_data['ONI']
        model = sm.OLS(y, sm.add_constant(X)).fit()
        y_pred = model.predict(sm.add_constant(X))
        fig.add_trace(go.Scatter(x=filtered_data['LON'], y=y_pred, mode='lines', name='Regression Line'))
        slope = model.params[1]
        slopes_text = f"All Years Slope: {slope:.4f}"
    else:
        slopes_text = "Insufficient data for regression"
    return fig, slopes_text

def generate_main_analysis(start_year, start_month, end_year, end_month, enso_phase, typhoon_search):
    start_date = datetime(start_year, start_month, 1)
    end_date = datetime(end_year, end_month, 28)
    filtered_data = merged_data[
        (merged_data['ISO_TIME'] >= start_date) & 
        (merged_data['ISO_TIME'] <= end_date)
    ]
    filtered_data['ENSO_Phase'] = filtered_data['ONI'].apply(classify_enso_phases)
    if enso_phase != 'all':
        filtered_data = filtered_data[filtered_data['ENSO_Phase'] == enso_phase.capitalize()]
    
    tracks_fig = generate_typhoon_tracks(filtered_data, typhoon_search)
    wind_scatter = generate_wind_oni_scatter(filtered_data, typhoon_search)
    pressure_scatter = generate_pressure_oni_scatter(filtered_data, typhoon_search)
    regression_fig, slopes_text = generate_regression_analysis(filtered_data)
    
    return tracks_fig, wind_scatter, pressure_scatter, regression_fig, slopes_text

# Path animation function
def generate_path_animation(year, typhoon, standard):
    typhoon_id = typhoon.split('(')[-1].strip(')')
    storm = ibtracs.get_storm(typhoon_id)
    fig = go.Figure()
    fig.add_trace(go.Scattergeo(lon=storm.lon, lat=storm.lat, mode='lines', line=dict(width=2, color='gray')))
    fig.add_trace(go.Scattergeo(lon=[storm.lon[0]], lat=[storm.lat[0]], mode='markers',
                                marker=dict(size=10, color='green', symbol='star'), name='Start'))
    frames = [
        go.Frame(data=[
            go.Scattergeo(lon=storm.lon[:i+1], lat=storm.lat[:i+1], mode='lines', line=dict(width=2, color='blue')),
            go.Scattergeo(lon=[storm.lon[i]], lat=[storm.lat[i]], mode='markers',
                          marker=dict(size=10, color=categorize_typhoon_by_standard(storm.vmax[i], standard)[1]))
        ], name=f"frame{i}") for i in range(len(storm.time))
    ]
    fig.frames = frames
    fig.update_layout(
        title=f"{year} {storm.name} Typhoon Path",
        geo=dict(projection_type='natural earth', showland=True),
        updatemenus=[{"buttons": [{"label": "Play", "method": "animate", "args": [None, {"frame": {"duration": 100}}]},
                                  {"label": "Pause", "method": "animate", "args": [[None], {"mode": "immediate"}]}]}]
    )
    return fig

def categorize_typhoon_by_standard(wind_speed, standard):
    if standard == 'taiwan':
        wind_speed_ms = wind_speed * 0.514444
        if wind_speed_ms >= 51.0:
            return 'Strong Typhoon', taiwan_standard['Strong Typhoon']['color']
        elif wind_speed_ms >= 33.7:
            return 'Medium Typhoon', taiwan_standard['Medium Typhoon']['color']
        elif wind_speed_ms >= 17.2:
            return 'Mild Typhoon', taiwan_standard['Mild Typhoon']['color']
        return 'Tropical Depression', taiwan_standard['Tropical Depression']['color']
    else:
        if wind_speed >= 137:
            return 'C5 Super Typhoon', atlantic_standard['C5 Super Typhoon']['color']
        elif wind_speed >= 113:
            return 'C4 Very Strong Typhoon', atlantic_standard['C4 Very Strong Typhoon']['color']
        elif wind_speed >= 96:
            return 'C3 Strong Typhoon', atlantic_standard['C3 Strong Typhoon']['color']
        elif wind_speed >= 83:
            return 'C2 Typhoon', atlantic_standard['C2 Typhoon']['color']
        elif wind_speed >= 64:
            return 'C1 Typhoon', atlantic_standard['C1 Typhoon']['color']
        elif wind_speed >= 34:
            return 'Tropical Storm', atlantic_standard['Tropical Storm']['color']
        return 'Tropical Depression', atlantic_standard['Tropical Depression']['color']

# Logistic regression functions
def perform_wind_regression(start_year, start_month, end_year, end_month):
    start_date = datetime(start_year, start_month, 1)
    end_date = datetime(end_year, end_month, 28)
    data = merged_data[(merged_data['ISO_TIME'] >= start_date) & (merged_data['ISO_TIME'] <= end_date)].dropna(subset=['USA_WIND', 'ONI'])
    data['severe_typhoon'] = (data['USA_WIND'] >= 64).astype(int)
    X = sm.add_constant(data['ONI'])
    y = data['severe_typhoon']
    model = sm.Logit(y, X).fit()
    beta_1, exp_beta_1, p_value = model.params['ONI'], np.exp(model.params['ONI']), model.pvalues['ONI']
    return f"Wind Regression: β1={beta_1:.4f}, Odds Ratio={exp_beta_1:.4f}, P-value={p_value:.4f}"

def perform_pressure_regression(start_year, start_month, end_year, end_month):
    start_date = datetime(start_year, start_month, 1)
    end_date = datetime(end_year, end_month, 28)
    data = merged_data[(merged_data['ISO_TIME'] >= start_date) & (merged_data['ISO_TIME'] <= end_date)].dropna(subset=['USA_PRES', 'ONI'])
    data['intense_typhoon'] = (data['USA_PRES'] <= 950).astype(int)
    X = sm.add_constant(data['ONI'])
    y = data['intense_typhoon']
    model = sm.Logit(y, X).fit()
    beta_1, exp_beta_1, p_value = model.params['ONI'], np.exp(model.params['ONI']), model.pvalues['ONI']
    return f"Pressure Regression: β1={beta_1:.4f}, Odds Ratio={exp_beta_1:.4f}, P-value={p_value:.4f}"

def perform_longitude_regression(start_year, start_month, end_year, end_month):
    start_date = datetime(start_year, start_month, 1)
    end_date = datetime(end_year, end_month, 28)
    data = merged_data[(merged_data['ISO_TIME'] >= start_date) & (merged_data['ISO_TIME'] <= end_date)].dropna(subset=['LON', 'ONI'])
    data['western_typhoon'] = (data['LON'] <= 140).astype(int)
    X = sm.add_constant(data['ONI'])
    y = data['western_typhoon']
    model = sm.Logit(y, X).fit()
    beta_1, exp_beta_1, p_value = model.params['ONI'], np.exp(model.params['ONI']), model.pvalues['ONI']
    return f"Longitude Regression: β1={beta_1:.4f}, Odds Ratio={exp_beta_1:.4f}, P-value={p_value:.4f}"

# Gradio interface
with gr.Blocks(title="Typhoon Analysis Dashboard") as demo:
    gr.Markdown("# Typhoon Analysis Dashboard")
    
    with gr.Tab("Overview"):
        gr.Markdown("""
        ## Welcome to the Typhoon Analysis Dashboard
        
        This dashboard allows you to analyze typhoon data in relation to ENSO phases.
        
        ### Features:
        - **Track Visualization**: View typhoon tracks by time period and ENSO phase
        - **Statistical Analysis**: Examine relationships between ONI values and typhoon characteristics
        - **Path Animation**: Watch animated typhoon paths with intensity classification
        - **Regression Analysis**: Perform statistical regression on typhoon data
        
        Select a tab above to begin your analysis.
        """)
    
    with gr.Tab("Track Visualization"):
        with gr.Row():
            start_year = gr.Number(label="Start Year", value=2000, minimum=1900, maximum=2024, step=1)
            start_month = gr.Dropdown(label="Start Month", choices=list(range(1, 13)), value=1)
            end_year = gr.Number(label="End Year", value=2024, minimum=1900, maximum=2024, step=1)
            end_month = gr.Dropdown(label="End Month", choices=list(range(1, 13)), value=6)
            enso_phase = gr.Dropdown(label="ENSO Phase", choices=['all', 'El Nino', 'La Nina', 'Neutral'], value='all')
            typhoon_search = gr.Textbox(label="Typhoon Search")
        analyze_btn = gr.Button("Generate Tracks")
        
        # Display all tracks
        tracks_plot = gr.Plot(label="Typhoon Tracks", elem_id="tracks_plot")
        typhoon_count = gr.Textbox(label="Number of Typhoons Displayed")
        
        # Enhanced function to show all track data
        def get_full_tracks(start_year, start_month, end_year, end_month, enso_phase, typhoon_search):
            start_date = datetime(start_year, start_month, 1)
            end_date = datetime(end_year, end_month, 28)
            filtered_data = merged_data[
                (merged_data['ISO_TIME'] >= start_date) & 
                (merged_data['ISO_TIME'] <= end_date)
            ]
            filtered_data['ENSO_Phase'] = filtered_data['ONI'].apply(classify_enso_phases)
            
            if enso_phase != 'all':
                filtered_data = filtered_data[filtered_data['ENSO_Phase'] == enso_phase.capitalize()]
            
            # Get all unique storms
            unique_storms = filtered_data['SID'].unique()
            count = len(unique_storms)
            
            # Create the map with all tracks
            fig = go.Figure()
            
            # Add all tracks
            for sid in unique_storms:
                storm_data = typhoon_data[typhoon_data['SID'] == sid]
                name = storm_data['NAME'].iloc[0] if not pd.isna(storm_data['NAME'].iloc[0]) else "Unnamed"
                
                # Get ENSO phase color
                storm_oni = filtered_data[filtered_data['SID'] == sid]['ONI'].iloc[0]
                color = 'red' if storm_oni >= 0.5 else ('blue' if storm_oni <= -0.5 else 'green')
                
                # Add the track line
                fig.add_trace(go.Scattergeo(
                    lon=storm_data['LON'], 
                    lat=storm_data['LAT'], 
                    mode='lines',
                    name=f"{name} ({storm_data['SEASON'].iloc[0]})",
                    line=dict(width=1.5, color=color),
                    hoverinfo="name"
                ))
            
            # Highlight searched typhoon if specified
            if typhoon_search:
                search_mask = typhoon_data['NAME'].str.contains(typhoon_search, case=False, na=False)
                if search_mask.any():
                    for sid in typhoon_data[search_mask]['SID'].unique():
                        storm_data = typhoon_data[typhoon_data['SID'] == sid]
                        fig.add_trace(go.Scattergeo(
                            lon=storm_data['LON'], 
                            lat=storm_data['LAT'], 
                            mode='lines+markers',
                            name=f"MATCHED: {storm_data['NAME'].iloc[0]} ({storm_data['SEASON'].iloc[0]})",
                            line=dict(width=3, color='yellow'),
                            marker=dict(size=5),
                            hoverinfo="name"
                        ))
            
            # Add colorbar/legend for ENSO phases
            fig.update_layout(
                title=f"Typhoon Tracks ({start_year}-{start_month} to {end_year}-{end_month})",
                geo=dict(
                    projection_type='natural earth',
                    showland=True,
                    showcoastlines=True,
                    landcolor='rgb(243, 243, 243)',
                    countrycolor='rgb(204, 204, 204)',
                    coastlinecolor='rgb(204, 204, 204)',
                    center=dict(lon=140, lat=20),
                    projection_scale=3
                ),
                legend_title="Typhoons by ENSO Phase",
                showlegend=True,
                height=700
            )
            
            # Add annotations explaining colors
            fig.add_annotation(
                x=0.02, y=0.98, xref="paper", yref="paper",
                text="Red: El Niño, Blue: La Niña, Green: Neutral",
                showarrow=False, align="left",
                bgcolor="rgba(255,255,255,0.8)"
            )
            
            return fig, f"Total typhoons displayed: {count}"
            
        analyze_btn.click(
            fn=get_full_tracks,
            inputs=[start_year, start_month, end_year, end_month, enso_phase, typhoon_search],
            outputs=[tracks_plot, typhoon_count]
        )
    
    with gr.Tab("Wind Analysis"):
        with gr.Row():
            wind_start_year = gr.Number(label="Start Year", value=2000, minimum=1900, maximum=2024, step=1)
            wind_start_month = gr.Dropdown(label="Start Month", choices=list(range(1, 13)), value=1)
            wind_end_year = gr.Number(label="End Year", value=2024, minimum=1900, maximum=2024, step=1)
            wind_end_month = gr.Dropdown(label="End Month", choices=list(range(1, 13)), value=6)
            wind_enso_phase = gr.Dropdown(label="ENSO Phase", choices=['all', 'El Nino', 'La Nina', 'Neutral'], value='all')
            wind_typhoon_search = gr.Textbox(label="Typhoon Search")
        wind_analyze_btn = gr.Button("Generate Wind Analysis")
        wind_scatter = gr.Plot(label="Wind Speed vs ONI")
        wind_regression_results = gr.Textbox(label="Wind Regression Results")
        
        # Fixed function for wind analysis
        def get_wind_analysis(start_year, start_month, end_year, end_month, enso_phase, typhoon_search):
            results = generate_main_analysis(start_year, start_month, end_year, end_month, enso_phase, typhoon_search)
            regression = perform_wind_regression(start_year, start_month, end_year, end_month)
            return results[1], regression
            
        wind_analyze_btn.click(
            fn=get_wind_analysis,
            inputs=[wind_start_year, wind_start_month, wind_end_year, wind_end_month, wind_enso_phase, wind_typhoon_search],
            outputs=[wind_scatter, wind_regression_results]
        )
    
    with gr.Tab("Pressure Analysis"):
        with gr.Row():
            pressure_start_year = gr.Number(label="Start Year", value=2000, minimum=1900, maximum=2024, step=1)
            pressure_start_month = gr.Dropdown(label="Start Month", choices=list(range(1, 13)), value=1)
            pressure_end_year = gr.Number(label="End Year", value=2024, minimum=1900, maximum=2024, step=1)
            pressure_end_month = gr.Dropdown(label="End Month", choices=list(range(1, 13)), value=6)
            pressure_enso_phase = gr.Dropdown(label="ENSO Phase", choices=['all', 'El Nino', 'La Nina', 'Neutral'], value='all')
            pressure_typhoon_search = gr.Textbox(label="Typhoon Search")
        pressure_analyze_btn = gr.Button("Generate Pressure Analysis")
        pressure_scatter = gr.Plot(label="Pressure vs ONI")
        pressure_regression_results = gr.Textbox(label="Pressure Regression Results")
        
        # Fixed function for pressure analysis
        def get_pressure_analysis(start_year, start_month, end_year, end_month, enso_phase, typhoon_search):
            results = generate_main_analysis(start_year, start_month, end_year, end_month, enso_phase, typhoon_search)
            regression = perform_pressure_regression(start_year, start_month, end_year, end_month)
            return results[2], regression
            
        pressure_analyze_btn.click(
            fn=get_pressure_analysis,
            inputs=[pressure_start_year, pressure_start_month, pressure_end_year, pressure_end_month, pressure_enso_phase, pressure_typhoon_search],
            outputs=[pressure_scatter, pressure_regression_results]
        )
    
    with gr.Tab("Longitude Analysis"):
        with gr.Row():
            lon_start_year = gr.Number(label="Start Year", value=2000, minimum=1900, maximum=2024, step=1)
            lon_start_month = gr.Dropdown(label="Start Month", choices=list(range(1, 13)), value=1)
            lon_end_year = gr.Number(label="End Year", value=2024, minimum=1900, maximum=2024, step=1)
            lon_end_month = gr.Dropdown(label="End Month", choices=list(range(1, 13)), value=6)
            lon_enso_phase = gr.Dropdown(label="ENSO Phase", choices=['all', 'El Nino', 'La Nina', 'Neutral'], value='all')
            lon_typhoon_search = gr.Textbox(label="Typhoon Search (Optional)")
        lon_analyze_btn = gr.Button("Generate Longitude Analysis")
        regression_plot = gr.Plot(label="Longitude vs ONI")
        slopes_text = gr.Textbox(label="Regression Slopes")
        lon_regression_results = gr.Textbox(label="Longitude Regression Results")
        
        # Fixed function for longitude analysis
        def get_longitude_analysis(start_year, start_month, end_year, end_month, enso_phase, typhoon_search):
            results = generate_main_analysis(start_year, start_month, end_year, end_month, enso_phase, typhoon_search)
            regression = perform_longitude_regression(start_year, start_month, end_year, end_month)
            return results[3], results[4], regression
            
        lon_analyze_btn.click(
            fn=get_longitude_analysis,
            inputs=[lon_start_year, lon_start_month, lon_end_year, lon_end_month, lon_enso_phase, lon_typhoon_search],
            outputs=[regression_plot, slopes_text, lon_regression_results]
        )
    
    with gr.Tab("Typhoon Path Animation"):
        with gr.Row():
            year_dropdown = gr.Dropdown(label="Year", choices=[str(y) for y in range(1950, 2025)], value="2024")
            typhoon_dropdown = gr.Dropdown(label="Typhoon")
            standard_dropdown = gr.Dropdown(label="Classification Standard", 
                                          choices=['atlantic', 'taiwan'], value='atlantic')
        
        # Completely redesigned animation function to show track movement clearly
        def generate_improved_animation(year, typhoon, standard):
            if not typhoon:
                return None
                
            typhoon_id = typhoon.split('(')[-1].strip(')')
            storm = ibtracs.get_storm(typhoon_id)
            
            # Create frames that show the growing track
            frames = []
            
            # Add frames showing growing track
            for i in range(len(storm.time)):
                # Get wind category and color
                category, color = categorize_typhoon_by_standard(storm.vmax[i], standard)
                
                # Create frame showing path up to current point
                frame_data = [
                    # Path line up to current point
                    go.Scattergeo(
                        lon=storm.lon[:i+1], 
                        lat=storm.lat[:i+1], 
                        mode='lines', 
                        line=dict(width=2, color='blue'),
                        name="Track"
                    ),
                    # Current position
                    go.Scattergeo(
                        lon=[storm.lon[i]], 
                        lat=[storm.lat[i]], 
                        mode='markers', 
                        marker=dict(size=12, color=color),
                        name=f"Current Position",
                        text=f"Time: {storm.time[i].strftime('%Y-%m-%d %H:%M')}<br>Wind: {storm.vmax[i]} kt<br>Category: {category}"
                    )
                ]
                
                # Add previous positions as smaller markers
                if i > 0:
                    frame_data.append(
                        go.Scattergeo(
                            lon=storm.lon[:i], 
                            lat=storm.lat[:i], 
                            mode='markers', 
                            marker=dict(size=5, color='rgba(100,100,100,0.5)'),
                            name="Previous Positions",
                            showlegend=False
                        )
                    )
                
                frames.append(go.Frame(data=frame_data, name=f"frame{i}"))
            
            # Initial figure showing start point
            fig = go.Figure(
                data=[
                    go.Scattergeo(
                        lon=[storm.lon[0]], 
                        lat=[storm.lat[0]],
                        mode='markers', 
                        marker=dict(size=12, color='green'),
                        name="Starting Position",
                        text=f"Start: {storm.time[0].strftime('%Y-%m-%d %H:%M')}"
                    )
                ],
                frames=frames
            )
            
            # Add category legend
            if standard == 'atlantic':
                for cat, details in atlantic_standard.items():
                    fig.add_trace(go.Scattergeo(
                        lon=[None], lat=[None], mode='markers',
                        marker=dict(size=10, color=details['color']),
                        name=cat
                    ))
            else:
                for cat, details in taiwan_standard.items():
                    fig.add_trace(go.Scattergeo(
                        lon=[None], lat=[None], mode='markers',
                        marker=dict(size=10, color=details['color']),
                        name=cat
                    ))
            
            # Focus map on storm area
            min_lat, max_lat = min(storm.lat), max(storm.lat)
            min_lon, max_lon = min(storm.lon), max(storm.lon)
            lat_padding = (max_lat - min_lat) * 0.3 or 5
            lon_padding = (max_lon - min_lon) * 0.3 or 5
            
            # Update layout with better animation controls
            fig.update_layout(
                title=f"{year} {storm.name} Typhoon Path",
                geo=dict(
                    projection_type='natural earth',
                    showland=True,
                    showcoastlines=True,
                    landcolor='rgb(243, 243, 243)',
                    countrycolor='rgb(204, 204, 204)',
                    coastlinecolor='rgb(204, 204, 204)',
                    showocean=True,
                    oceancolor='rgb(230, 230, 255)',
                    lataxis={'range': [min_lat - lat_padding, max_lat + lat_padding]},
                    lonaxis={'range': [min_lon - lon_padding, max_lon + lon_padding]}
                ),
                updatemenus=[{
                    "buttons": [
                        {
                            "args": [None, {"frame": {"duration": 100, "redraw": True}, "fromcurrent": True, "mode": "immediate"}],
                            "label": "Play",
                            "method": "animate"
                        },
                        {
                            "args": [[None], {"frame": {"duration": 0, "redraw": True}, "mode": "immediate"}],
                            "label": "Pause",
                            "method": "animate"
                        }
                    ],
                    "direction": "left",
                    "pad": {"r": 10, "t": 10},
                    "type": "buttons",
                    "x": 0.1,
                    "y": 0
                }],
                sliders=[{
                    "active": 0,
                    "yanchor": "top",
                    "xanchor": "left",
                    "currentvalue": {
                        "font": {"size": 12},
                        "prefix": "Time: ",
                        "visible": True,
                        "xanchor": "right"
                    },
                    "pad": {"b": 10, "t": 50},
                    "len": 0.9,
                    "x": 0.1,
                    "y": 0,
                    "steps": [
                        {
                            "args": [[f.name], {
                                "frame": {"duration": 0, "redraw": True},
                                "mode": "immediate"
                            }],
                            "label": storm.time[i].strftime('%m/%d %H:00') if i < len(storm.time) else "",
                            "method": "animate"
                        } for i, f in enumerate(frames)
                    ]
                }],
                height=700,
                showlegend=True,
                legend=dict(
                    yanchor="top",
                    y=0.99,
                    xanchor="left",
                    x=0.01,
                    bgcolor="rgba(255, 255, 255, 0.8)"
                )
            )
            
            return fig
        
        animate_btn = gr.Button("Generate Animation")
        path_plot = gr.Plot(label="Typhoon Path Animation", elem_id="animation_plot")
        animation_info = gr.Markdown("""
        ### Animation Instructions
        1. Select a year and typhoon from the dropdowns
        2. Click "Generate Animation"
        3. Use the play button to start the animation
        4. Use the slider to manually control the animation timeline
        5. The animation shows the typhoon track developing over time, with the current position highlighted
        6. Colors indicate typhoon intensity according to the selected classification standard
        """)
        
        # Year dropdown change function
        def update_typhoon_options(year):
            season = ibtracs.get_season(int(year))
            storm_summary = season.summary()
            options = [f"{storm_summary['name'][i]} ({storm_summary['id'][i]})" for i in range(storm_summary['season_storms'])]
            return gr.update(choices=options, value=options[0] if options else None)
        
        year_dropdown.change(fn=update_typhoon_options, inputs=year_dropdown, outputs=typhoon_dropdown)
        animate_btn.click(
            fn=generate_improved_animation,
            inputs=[year_dropdown, typhoon_dropdown, standard_dropdown],
            outputs=path_plot
        )

    # Custom CSS for better spacing and to ensure plots are visible
    gr.HTML("""
    <style>
    #tracks_plot, #animation_plot {
        height: 700px !important;
    }
    .plot-container {
        min-height: 600px;
    }
    </style>
    """)

demo.launch()