import pandas as pd
import numpy as np
from datetime import datetime
from data import extract_model_data
import gradio as gr

def get_time_series_summary_dfs(historical_df: pd.DataFrame) -> dict:
    """Return dataframes for historical summary plots (failure rates, AMD tests, NVIDIA tests)."""

    # Group by date to get daily statistics
    daily_stats = []
    dates = sorted(historical_df['date'].unique())
    for date in dates:
        date_data = historical_df[historical_df['date'] == date]
        amd_passed = date_data['success_amd'].sum() if 'success_amd' in date_data.columns else 0
        amd_failed = (date_data['failed_multi_no_amd'].sum() + date_data['failed_single_no_amd'].sum()) if 'failed_multi_no_amd' in date_data.columns else 0
        amd_skipped = date_data['skipped_amd'].sum() if 'skipped_amd' in date_data.columns else 0
        amd_total = amd_passed + amd_failed + amd_skipped
        amd_failure_rate = (amd_failed / amd_total * 100) if amd_total > 0 else 0

        nvidia_passed = date_data['success_nvidia'].sum() if 'success_nvidia' in date_data.columns else 0
        nvidia_failed = (date_data['failed_multi_no_nvidia'].sum() + date_data['failed_single_no_nvidia'].sum()) if 'failed_multi_no_nvidia' in date_data.columns else 0
        nvidia_skipped = date_data['skipped_nvidia'].sum() if 'skipped_nvidia' in date_data.columns else 0
        nvidia_total = nvidia_passed + nvidia_failed + nvidia_skipped
        nvidia_failure_rate = (nvidia_failed / nvidia_total * 100) if nvidia_total > 0 else 0

        daily_stats.append({
            'date': date,
            'amd_failure_rate': amd_failure_rate,
            'nvidia_failure_rate': nvidia_failure_rate,
            'amd_passed': amd_passed,
            'amd_failed': amd_failed,
            'amd_skipped': amd_skipped,
            'nvidia_passed': nvidia_passed,
            'nvidia_failed': nvidia_failed,
            'nvidia_skipped': nvidia_skipped
        })

    # Failure rate dataframe
    failure_rate_data = []
    for i, stat in enumerate(daily_stats):
        amd_change = stat['amd_failure_rate'] - daily_stats[i-1]['amd_failure_rate'] if i > 0 else 0
        nvidia_change = stat['nvidia_failure_rate'] - daily_stats[i-1]['nvidia_failure_rate'] if i > 0 else 0
        failure_rate_data.extend([
            {'date': stat['date'], 'failure_rate': stat['amd_failure_rate'], 'platform': 'AMD', 'change': amd_change},
            {'date': stat['date'], 'failure_rate': stat['nvidia_failure_rate'], 'platform': 'NVIDIA', 'change': nvidia_change}
        ])
    failure_rate_df = pd.DataFrame(failure_rate_data)

    # AMD tests dataframe
    amd_data = []
    for i, stat in enumerate(daily_stats):
        passed_change = stat['amd_passed'] - daily_stats[i-1]['amd_passed'] if i > 0 else 0
        failed_change = stat['amd_failed'] - daily_stats[i-1]['amd_failed'] if i > 0 else 0
        skipped_change = stat['amd_skipped'] - daily_stats[i-1]['amd_skipped'] if i > 0 else 0
        amd_data.extend([
            {'date': stat['date'], 'count': stat['amd_passed'], 'test_type': 'Passed', 'change': passed_change},
            {'date': stat['date'], 'count': stat['amd_failed'], 'test_type': 'Failed', 'change': failed_change},
            {'date': stat['date'], 'count': stat['amd_skipped'], 'test_type': 'Skipped', 'change': skipped_change}
        ])
    amd_df = pd.DataFrame(amd_data)

    # NVIDIA tests dataframe
    nvidia_data = []
    for i, stat in enumerate(daily_stats):
        passed_change = stat['nvidia_passed'] - daily_stats[i-1]['nvidia_passed'] if i > 0 else 0
        failed_change = stat['nvidia_failed'] - daily_stats[i-1]['nvidia_failed'] if i > 0 else 0
        skipped_change = stat['nvidia_skipped'] - daily_stats[i-1]['nvidia_skipped'] if i > 0 else 0
        nvidia_data.extend([
            {'date': stat['date'], 'count': stat['nvidia_passed'], 'test_type': 'Passed', 'change': passed_change},
            {'date': stat['date'], 'count': stat['nvidia_failed'], 'test_type': 'Failed', 'change': failed_change},
            {'date': stat['date'], 'count': stat['nvidia_skipped'], 'test_type': 'Skipped', 'change': skipped_change}
        ])
    nvidia_df = pd.DataFrame(nvidia_data)

    return {
        'failure_rates_df': failure_rate_df,
        'amd_tests_df': amd_df,
        'nvidia_tests_df': nvidia_df,
    }

def get_model_time_series_dfs(historical_df: pd.DataFrame, model_name: str) -> dict:
    """Return dataframes for a specific model's historical plots (AMD, NVIDIA)."""
    model_data = historical_df[historical_df.index.str.lower() == model_name.lower()]
    
    if model_data.empty:
        empty_df = pd.DataFrame({'date': [], 'count': [], 'test_type': [], 'change': []})
        return {'amd_df': empty_df.copy(), 'nvidia_df': empty_df.copy()}

    dates = sorted(model_data['date'].unique())
    amd_data = []
    nvidia_data = []
    for i, date in enumerate(dates):
        date_data = model_data[model_data['date'] == date]
        row = date_data.iloc[0]

        amd_passed = row.get('success_amd', 0)
        amd_failed = row.get('failed_multi_no_amd', 0) + row.get('failed_single_no_amd', 0)
        amd_skipped = row.get('skipped_amd', 0)
        prev_row = model_data[model_data['date'] == dates[i-1]].iloc[0] if i > 0 and not model_data[model_data['date'] == dates[i-1]].empty else None
        amd_passed_change = amd_passed - (prev_row.get('success_amd', 0) if prev_row is not None else 0)
        amd_failed_change = amd_failed - (prev_row.get('failed_multi_no_amd', 0) + prev_row.get('failed_single_no_amd', 0) if prev_row is not None else 0)
        amd_skipped_change = amd_skipped - (prev_row.get('skipped_amd', 0) if prev_row is not None else 0)
        amd_data.extend([
            {'date': date, 'count': amd_passed, 'test_type': 'Passed', 'change': amd_passed_change},
            {'date': date, 'count': amd_failed, 'test_type': 'Failed', 'change': amd_failed_change},
            {'date': date, 'count': amd_skipped, 'test_type': 'Skipped', 'change': amd_skipped_change}
        ])

        nvidia_passed = row.get('success_nvidia', 0)
        nvidia_failed = row.get('failed_multi_no_nvidia', 0) + row.get('failed_single_no_nvidia', 0)
        nvidia_skipped = row.get('skipped_nvidia', 0)
        if prev_row is not None:
            prev_nvidia_passed = prev_row.get('success_nvidia', 0)
            prev_nvidia_failed = prev_row.get('failed_multi_no_nvidia', 0) + prev_row.get('failed_single_no_nvidia', 0)
            prev_nvidia_skipped = prev_row.get('skipped_nvidia', 0)
        else:
            prev_nvidia_passed = prev_nvidia_failed = prev_nvidia_skipped = 0
        nvidia_data.extend([
            {'date': date, 'count': nvidia_passed, 'test_type': 'Passed', 'change': nvidia_passed - prev_nvidia_passed},
            {'date': date, 'count': nvidia_failed, 'test_type': 'Failed', 'change': nvidia_failed - prev_nvidia_failed},
            {'date': date, 'count': nvidia_skipped, 'test_type': 'Skipped', 'change': nvidia_skipped - prev_nvidia_skipped}
        ])

    return {'amd_df': pd.DataFrame(amd_data), 'nvidia_df': pd.DataFrame(nvidia_data)}

def create_time_series_summary_gradio(historical_df: pd.DataFrame) -> dict:
    """Create time-series visualization for overall failure rates over time using Gradio native plots."""
    if historical_df.empty or 'date' not in historical_df.columns:
        # Return empty plots
        empty_df = pd.DataFrame({'date': [], 'failure_rate': [], 'platform': []})
        return {
            'failure_rates': gr.LinePlot(empty_df, x="date", y="failure_rate", color="platform", title="No historical data available", tooltip=["failure_rate", "date", "change"]),
            'amd_tests': gr.LinePlot(empty_df, x="date", y="failure_rate", color="platform", title="No historical data available", tooltip=["count", "date", "change"]),
            'nvidia_tests': gr.LinePlot(empty_df, x="date", y="failure_rate", color="platform", title="No historical data available", tooltip=["count", "date", "change"])
        }
    
    # Group by date to get daily statistics
    daily_stats = []
    dates = sorted(historical_df['date'].unique())
    
    for date in dates:
        date_data = historical_df[historical_df['date'] == date]
        
        # Calculate AMD stats - use the correct column names from the data structure
        amd_passed = date_data['success_amd'].sum() if 'success_amd' in date_data.columns else 0
        amd_failed = (date_data['failed_multi_no_amd'].sum() + date_data['failed_single_no_amd'].sum()) if 'failed_multi_no_amd' in date_data.columns else 0
        amd_skipped = date_data['skipped_amd'].sum() if 'skipped_amd' in date_data.columns else 0
        amd_total = amd_passed + amd_failed + amd_skipped
        amd_failure_rate = (amd_failed / amd_total * 100) if amd_total > 0 else 0
        
        # Calculate NVIDIA stats - use the correct column names from the data structure
        nvidia_passed = date_data['success_nvidia'].sum() if 'success_nvidia' in date_data.columns else 0
        nvidia_failed = (date_data['failed_multi_no_nvidia'].sum() + date_data['failed_single_no_nvidia'].sum()) if 'failed_multi_no_nvidia' in date_data.columns else 0
        nvidia_skipped = date_data['skipped_nvidia'].sum() if 'skipped_nvidia' in date_data.columns else 0
        nvidia_total = nvidia_passed + nvidia_failed + nvidia_skipped
        nvidia_failure_rate = (nvidia_failed / nvidia_total * 100) if nvidia_total > 0 else 0
        
        daily_stats.append({
            'date': date,
            'amd_failure_rate': amd_failure_rate,
            'nvidia_failure_rate': nvidia_failure_rate,
            'amd_passed': amd_passed,
            'amd_failed': amd_failed,
            'amd_skipped': amd_skipped,
            'nvidia_passed': nvidia_passed,
            'nvidia_failed': nvidia_failed,
            'nvidia_skipped': nvidia_skipped
        })
    
    # Create failure rate data
    failure_rate_data = []
    for i, stat in enumerate(daily_stats):
        # Calculate change from previous point
        amd_change = 0
        nvidia_change = 0
        if i > 0:
            amd_change = stat['amd_failure_rate'] - daily_stats[i-1]['amd_failure_rate']
            nvidia_change = stat['nvidia_failure_rate'] - daily_stats[i-1]['nvidia_failure_rate']
        
        failure_rate_data.extend([
            {'date': stat['date'], 'failure_rate': stat['amd_failure_rate'], 'platform': 'AMD', 'change': amd_change},
            {'date': stat['date'], 'failure_rate': stat['nvidia_failure_rate'], 'platform': 'NVIDIA', 'change': nvidia_change}
        ])
    
    failure_rate_df = pd.DataFrame(failure_rate_data)
    
    # Create AMD test results data
    amd_data = []
    for i, stat in enumerate(daily_stats):
        # Calculate change from previous point for each test type
        passed_change = 0
        failed_change = 0
        skipped_change = 0
        if i > 0:
            passed_change = stat['amd_passed'] - daily_stats[i-1]['amd_passed']
            failed_change = stat['amd_failed'] - daily_stats[i-1]['amd_failed']
            skipped_change = stat['amd_skipped'] - daily_stats[i-1]['amd_skipped']
        
        amd_data.extend([
            {'date': stat['date'], 'count': stat['amd_passed'], 'test_type': 'Passed', 'change': passed_change},
            {'date': stat['date'], 'count': stat['amd_failed'], 'test_type': 'Failed', 'change': failed_change},
            {'date': stat['date'], 'count': stat['amd_skipped'], 'test_type': 'Skipped', 'change': skipped_change}
        ])
    
    amd_df = pd.DataFrame(amd_data)
    
    # Create NVIDIA test results data
    nvidia_data = []
    for i, stat in enumerate(daily_stats):
        # Calculate change from previous point for each test type
        passed_change = 0
        failed_change = 0
        skipped_change = 0
        if i > 0:
            passed_change = stat['nvidia_passed'] - daily_stats[i-1]['nvidia_passed']
            failed_change = stat['nvidia_failed'] - daily_stats[i-1]['nvidia_failed']
            skipped_change = stat['nvidia_skipped'] - daily_stats[i-1]['nvidia_skipped']
        
        nvidia_data.extend([
            {'date': stat['date'], 'count': stat['nvidia_passed'], 'test_type': 'Passed', 'change': passed_change},
            {'date': stat['date'], 'count': stat['nvidia_failed'], 'test_type': 'Failed', 'change': failed_change},
            {'date': stat['date'], 'count': stat['nvidia_skipped'], 'test_type': 'Skipped', 'change': skipped_change}
        ])
    
    nvidia_df = pd.DataFrame(nvidia_data)
    
    return {
        'failure_rates': gr.LinePlot(
            failure_rate_df,
            x="date",
            y="failure_rate",
            color="platform",
            color_map={"AMD": "#FF6B6B", "NVIDIA": "#4ECDC4"},
            title="Overall Failure Rates Over Time",
            tooltip=["failure_rate", "date", "change"],
            height=300,
            x_label_angle=45,
            y_title="Failure Rate (%)"
        ),
        'amd_tests': gr.LinePlot(
            amd_df,
            x="date",
            y="count",
            color="test_type",
            color_map={"Passed": "#4CAF50", "Failed": "#E53E3E", "Skipped": "#FFA500"},
            title="AMD Test Results Over Time",
            tooltip=["count", "date", "change"],
            height=300,
            x_label_angle=45,
            y_title="Number of Tests"
        ),
        'nvidia_tests': gr.LinePlot(
            nvidia_df,
            x="date",
            y="count",
            color="test_type",
            color_map={"Passed": "#4CAF50", "Failed": "#E53E3E", "Skipped": "#FFA500"},
            title="NVIDIA Test Results Over Time",
            tooltip=["count", "date", "change"],
            height=300,
            x_label_angle=45,
            y_title="Number of Tests"
        )
    }


def create_model_time_series_gradio(historical_df: pd.DataFrame, model_name: str) -> dict:
    """Create time-series visualization for a specific model using Gradio native plots."""
    if historical_df.empty or 'date' not in historical_df.columns:
        # Return empty plots
        empty_df = pd.DataFrame({'date': [], 'count': [], 'test_type': []})
        return {
            'amd_plot': gr.LinePlot(empty_df, x="date", y="count", color="test_type", title=f"{model_name.upper()} - AMD Results Over Time", tooltip=["count", "date", "change"]),
            'nvidia_plot': gr.LinePlot(empty_df, x="date", y="count", color="test_type", title=f"{model_name.upper()} - NVIDIA Results Over Time", tooltip=["count", "date", "change"])
        }
    
    # Filter data for the specific model (model_name is the index)
    model_data = historical_df[historical_df.index.str.lower() == model_name.lower()]
    
    if model_data.empty:
        # Return empty plots
        empty_df = pd.DataFrame({'date': [], 'count': [], 'test_type': []})
        return {
            'amd_plot': gr.LinePlot(empty_df, x="date", y="count", color="test_type", title=f"{model_name.upper()} - AMD Results Over Time", tooltip=["count", "date", "change"]),
            'nvidia_plot': gr.LinePlot(empty_df, x="date", y="count", color="test_type", title=f"{model_name.upper()} - NVIDIA Results Over Time", tooltip=["count", "date", "change"])
        }
    
    # Group by date
    dates = sorted(model_data['date'].unique())
    
    amd_data = []
    nvidia_data = []
    
    for i, date in enumerate(dates):
        date_data = model_data[model_data['date'] == date]
        
        if not date_data.empty:
            # Get the first row for this date (should be only one)
            row = date_data.iloc[0]
            
            # AMD data - use the correct column names from the data structure
            amd_passed = row.get('success_amd', 0)
            amd_failed = row.get('failed_multi_no_amd', 0) + row.get('failed_single_no_amd', 0)
            amd_skipped = row.get('skipped_amd', 0)
            
            # Calculate change from previous point
            passed_change = 0
            failed_change = 0
            skipped_change = 0
            if i > 0:
                prev_date_data = model_data[model_data['date'] == dates[i-1]]
                if not prev_date_data.empty:
                    prev_row = prev_date_data.iloc[0]
                    prev_amd_passed = prev_row.get('success_amd', 0)
                    prev_amd_failed = prev_row.get('failed_multi_no_amd', 0) + prev_row.get('failed_single_no_amd', 0)
                    prev_amd_skipped = prev_row.get('skipped_amd', 0)
                    
                    passed_change = amd_passed - prev_amd_passed
                    failed_change = amd_failed - prev_amd_failed
                    skipped_change = amd_skipped - prev_amd_skipped
            
            amd_data.extend([
                {'date': date, 'count': amd_passed, 'test_type': 'Passed', 'change': passed_change},
                {'date': date, 'count': amd_failed, 'test_type': 'Failed', 'change': failed_change},
                {'date': date, 'count': amd_skipped, 'test_type': 'Skipped', 'change': skipped_change}
            ])
            
            # NVIDIA data - use the correct column names from the data structure
            nvidia_passed = row.get('success_nvidia', 0)
            nvidia_failed = row.get('failed_multi_no_nvidia', 0) + row.get('failed_single_no_nvidia', 0)
            nvidia_skipped = row.get('skipped_nvidia', 0)
            
            # Calculate change from previous point for NVIDIA
            nvidia_passed_change = 0
            nvidia_failed_change = 0
            nvidia_skipped_change = 0
            if i > 0:
                prev_date_data = model_data[model_data['date'] == dates[i-1]]
                if not prev_date_data.empty:
                    prev_row = prev_date_data.iloc[0]
                    prev_nvidia_passed = prev_row.get('success_nvidia', 0)
                    prev_nvidia_failed = prev_row.get('failed_multi_no_nvidia', 0) + prev_row.get('failed_single_no_nvidia', 0)
                    prev_nvidia_skipped = prev_row.get('skipped_nvidia', 0)
                    
                    nvidia_passed_change = nvidia_passed - prev_nvidia_passed
                    nvidia_failed_change = nvidia_failed - prev_nvidia_failed
                    nvidia_skipped_change = nvidia_skipped - prev_nvidia_skipped
            
            nvidia_data.extend([
                {'date': date, 'count': nvidia_passed, 'test_type': 'Passed', 'change': nvidia_passed_change},
                {'date': date, 'count': nvidia_failed, 'test_type': 'Failed', 'change': nvidia_failed_change},
                {'date': date, 'count': nvidia_skipped, 'test_type': 'Skipped', 'change': nvidia_skipped_change}
            ])
    
    amd_df = pd.DataFrame(amd_data)
    nvidia_df = pd.DataFrame(nvidia_data)
    
    return {
        'amd_plot': gr.LinePlot(
            amd_df,
            x="date",
            y="count",
            color="test_type",
            color_map={"Passed": "#4CAF50", "Failed": "#E53E3E", "Skipped": "#FFA500"},
            title=f"{model_name.upper()} - AMD Results Over Time",
            x_label_angle=45,
            y_title="Number of Tests",
            height=300,
            tooltip=["count", "date", "change"]
        ),
        'nvidia_plot': gr.LinePlot(
            nvidia_df,
            x="date",
            y="count",
            color="test_type",
            color_map={"Passed": "#4CAF50", "Failed": "#E53E3E", "Skipped": "#FFA500"},
            title=f"{model_name.upper()} - NVIDIA Results Over Time",
            x_label_angle=45,
            y_title="Number of Tests",
            height=300,
            tooltip=["count", "date", "change"]
        )
    }