# app.py
import os
os.environ['TF_ENABLE_ONEDNN_OPTS'] = '0'  # Suppress TensorFlow warnings

import gradio as gr
import pandas as pd
import numpy as np
from scipy.fft import fft, fftfreq
from sklearn.preprocessing import MinMaxScaler
from tensorflow.keras.models import Sequential
from tensorflow.keras.layers import LSTM, Dense
import matplotlib.pyplot as plt

def load_data(input_source):
    """Handle both uploaded files and URLs"""
    try:
        if isinstance(input_source, str) and input_source.startswith("http"):
            # Load from URL
            df = pd.read_csv(
                input_source,
                engine='python',
                on_bad_lines='warn',
                encoding='utf-8'
            )
        else:
            # Load from uploaded file
            df = pd.read_csv(
                input_source.name,
                engine='python',
                on_bad_lines='warn',
                encoding='utf-8'
            )
        
        # Common cleaning steps
        df = df.drop(columns=['Province/State', 'Lat', 'Long'], errors='ignore')
        df = df.groupby('Country/Region').sum().T
        df.index = pd.to_datetime(df.index)
        df['Global'] = df.sum(axis=1)
        return df['Global'].diff().fillna(0)
    
    except Exception as e:
        raise ValueError(f"Data loading failed: {str(e)}")

def analyze_data(input_source):
    try:
        if not input_source:
            return "⚠️ Please upload a file or enter a URL", None
        
        data = load_data(input_source)
        
        # Cycle detection
        N = len(data)
        yf = fft(data.values)
        xf = fftfreq(N, 1)[:N//2]
        cycle_days = int(1/xf[np.argmax(np.abs(yf[0:N//2]))])
        
        # Create plot
        fig, ax = plt.subplots(figsize=(10, 4))
        ax.plot(data.index, data.values, label='Daily Cases')
        ax.set_title("COVID-19 Analysis")
        ax.set_xlabel("Date")
        ax.set_ylabel("New Cases")
        ax.grid(True)
        plt.tight_layout()
        
        # Generate insights
        latest_avg = data[-30:].mean()
        trend = "↑ Rising" if data[-1] > data[-7] else "↓ Falling"
        
        return (
            f"🔍 Analysis Results:\n"
            f"- Dominant Cycle: {cycle_days} days\n"
            f"- 30-Day Average: {latest_avg:.1f} cases/day\n"
            f"- Current Trend: {trend}\n"
            f"✅ Analysis completed successfully!",
            fig
        )
        
    except Exception as e:
        return f"❌ Error: {str(e)}", None

# Create interface
with gr.Blocks(theme=gr.themes.Soft()) as app:
    gr.Markdown("# 🦠 COVID-19 Analysis Bot")
    gr.Markdown("Analyze case data from URLs or uploaded CSV files")
    
    with gr.Row():
        with gr.Column():
            file_upload = gr.File(label="1. Upload CSV", file_count='single')
            url_input = gr.Textbox(label="2. Or paste data URL here")
            submit_btn = gr.Button("Analyze ➔")
            
        with gr.Column():
            chat_output = gr.Chatbot(label="Analysis Results", height=300)
            plot_output = gr.Plot(label="Case Trend")
    
    # Link components
    submit_btn.click(
        fn=analyze_data,
        inputs=[gr.combine(file_upload, url_input)],
        outputs=[chat_output, plot_output]
    )
    
    # Examples
    gr.Examples(
        examples=[
            ["https://raw.githubusercontent.com/CSSEGISandData/COVID-19/master/data/time_series_covid19_confirmed_global.csv"]
        ],
        inputs=[url_input],
        label="Try this example URL:"
    )

if __name__ == "__main__":
    app.launch()