import pandas as pd
import numpy as np
from datetime import datetime
from data import extract_model_data
from utils import COLORS
import gradio as gr
import plotly.express as px
import plotly.graph_objects as go

def create_time_series_summary_gradio(historical_df: pd.DataFrame) -> dict:
    empty_fig = lambda title: go.Figure().update_layout(title=title, height=500, 
        font=dict(size=16, color='#CCCCCC'), paper_bgcolor='#000000', 
        plot_bgcolor='#1a1a1a', margin=dict(b=130)) or go.Figure()
    
    if historical_df.empty or 'date' not in historical_df.columns:
        ef = empty_fig("No historical data available")
        return {'failure_rates': ef, 'amd_tests': ef, 'nvidia_tests': ef}
    
    daily_stats = []
    for date in sorted(historical_df['date'].unique()):
        dd = historical_df[historical_df['date'] == date]
        counts = {'date': date}
        
        for platform in ['amd', 'nvidia']:
            tot_tests = tot_fails = p = f = s = 0
            for _, row in dd.iterrows():
                stats = extract_model_data(row)[0 if platform == 'amd' else 1]
                tot = stats['passed'] + stats['failed'] + stats['error']
                if tot > 0:
                    tot_tests += tot
                    tot_fails += stats['failed'] + stats['error']
                p += stats['passed']
                f += stats['failed'] + stats['error']
                s += stats['skipped']
            
            counts.update({f'{platform}_failure_rate': (tot_fails / tot_tests * 100) if tot_tests > 0 else 0,
                          f'{platform}_passed': p, f'{platform}_failed': f, f'{platform}_skipped': s})
        daily_stats.append(counts)
    
    fr_data = []
    for i, s in enumerate(daily_stats):
        for p in ['amd', 'nvidia']:
            chg = s[f'{p}_failure_rate'] - daily_stats[i-1][f'{p}_failure_rate'] if i > 0 else 0
            fr_data.append({'date': s['date'], 'failure_rate': s[f'{p}_failure_rate'], 
                           'platform': p.upper(), 'change': chg})
    
    def build_test_data(platform):
        data = []
        for i, s in enumerate(daily_stats):
            for tt in ['passed', 'failed', 'skipped']:
                chg = s[f'{platform}_{tt}'] - daily_stats[i-1][f'{platform}_{tt}'] if i > 0 else 0
                data.append({'date': s['date'], 'count': s[f'{platform}_{tt}'], 
                           'test_type': tt.capitalize(), 'change': chg})
        return pd.DataFrame(data)
    
    fr_df = pd.DataFrame(fr_data)
    
    fig_fr = go.Figure()
    for p, lc, mc in [('NVIDIA', '#76B900', '#FFFFFF'), ('AMD', '#ED1C24', '#404040')]:
        d = fr_df[fr_df['platform'] == p]
        if not d.empty:
            fig_fr.add_trace(go.Scatter(x=d['date'], y=d['failure_rate'], mode='lines+markers',
                name=p, line=dict(color=lc, width=3), 
                marker=dict(size=12, color=mc, line=dict(color=lc, width=2)),
                hovertemplate=f'<b>{p}</b><br>Date: %{{x}}<br>Failure Rate: %{{y:.2f}}%<extra></extra>'))
    
    fig_fr.update_layout(title="Overall Failure Rates Over Time", height=500,
        font=dict(size=16, color='#CCCCCC'), paper_bgcolor='#000000', plot_bgcolor='#1a1a1a',
        title_font_size=20, legend=dict(font=dict(size=16), bgcolor='rgba(0,0,0,0.5)', 
        orientation="h", yanchor="bottom", y=-0.4, xanchor="center", x=0.5),
        xaxis=dict(title='Date', title_font_size=16, tickfont_size=14, gridcolor='#333333', showgrid=True),
        yaxis=dict(title='Failure Rate (%)', title_font_size=16, tickfont_size=14, gridcolor='#333333', showgrid=True),
        hovermode='x unified', margin=dict(b=130))
    
    def create_line_fig(df, title):
        fig = px.line(df, x='date', y='count', color='test_type',
            color_discrete_map={"Passed": COLORS['passed'], "Failed": COLORS['failed'], "Skipped": COLORS['skipped']},
            title=title, labels={'count': 'Number of Tests', 'date': 'Date', 'test_type': 'Test Type'})
        fig.update_traces(mode='lines+markers', marker=dict(size=8), line=dict(width=3))
        fig.update_layout(height=500, font=dict(size=16, color='#CCCCCC'), paper_bgcolor='#000000',
            plot_bgcolor='#1a1a1a', title_font_size=20, legend=dict(font=dict(size=16), 
            bgcolor='rgba(0,0,0,0.5)', orientation="h", yanchor="bottom", y=-0.4, xanchor="center", x=0.5),
            xaxis=dict(title_font_size=16, tickfont_size=14, gridcolor='#333333', showgrid=True),
            yaxis=dict(title_font_size=16, tickfont_size=14, gridcolor='#333333', showgrid=True),
            hovermode='x unified', margin=dict(b=130))
        return fig
    
    return {'failure_rates': fig_fr, 
            'amd_tests': create_line_fig(build_test_data('amd'), "AMD Test Results Over Time"),
            'nvidia_tests': create_line_fig(build_test_data('nvidia'), "NVIDIA Test Results Over Time")}

def create_model_time_series_gradio(historical_df: pd.DataFrame, model_name: str) -> dict:
    def empty_figs():
        ef = lambda plat: go.Figure().update_layout(title=f"{model_name.upper()} - {plat} Results Over Time",
            height=500, font=dict(size=16, color='#CCCCCC'), paper_bgcolor='#000000',
            plot_bgcolor='#1a1a1a', margin=dict(b=130)) or go.Figure()
        return {'amd_plot': ef('AMD'), 'nvidia_plot': ef('NVIDIA')}
    
    if historical_df.empty or 'date' not in historical_df.columns:
        return empty_figs()
    
    md = historical_df[historical_df.index.str.lower() == model_name.lower()]
    if md.empty:
        return empty_figs()
    
    dates = sorted(md['date'].unique())
    
    def build_data(platform):
        data = []
        for i, date in enumerate(dates):
            dd = md[md['date'] == date]
            if dd.empty:
                continue
            r = dd.iloc[0]
            passed = r.get(f'success_{platform}', 0)
            failed = r.get(f'failed_multi_no_{platform}', 0) + r.get(f'failed_single_no_{platform}', 0)
            skipped = r.get(f'skipped_{platform}', 0)
            
            pc = fc = sc = 0
            if i > 0:
                prev_dd = md[md['date'] == dates[i-1]]
                if not prev_dd.empty:
                    pr = prev_dd.iloc[0]
                    pc = pr.get(f'success_{platform}', 0)
                    fc = pr.get(f'failed_multi_no_{platform}', 0) + pr.get(f'failed_single_no_{platform}', 0)
                    sc = pr.get(f'skipped_{platform}', 0)
            
            data.extend([
                {'date': date, 'count': passed, 'test_type': 'Passed', 'change': passed - pc},
                {'date': date, 'count': failed, 'test_type': 'Failed', 'change': failed - fc},
                {'date': date, 'count': skipped, 'test_type': 'Skipped', 'change': skipped - sc}
            ])
        return pd.DataFrame(data)
    
    def create_fig(df, platform):
        fig = px.line(df, x='date', y='count', color='test_type',
            color_discrete_map={"Passed": COLORS['passed'], "Failed": COLORS['failed'], "Skipped": COLORS['skipped']},
            title=f"{model_name.upper()} - {platform} Results Over Time",
            labels={'count': 'Number of Tests', 'date': 'Date', 'test_type': 'Test Type'})
        fig.update_traces(mode='lines+markers', marker=dict(size=8), line=dict(width=3))
        fig.update_layout(height=500, font=dict(size=16, color='#CCCCCC'), paper_bgcolor='#000000',
            plot_bgcolor='#1a1a1a', title_font_size=20, legend=dict(font=dict(size=16),
            bgcolor='rgba(0,0,0,0.5)', orientation="h", yanchor="bottom", y=-0.4, xanchor="center", x=0.5),
            xaxis=dict(title_font_size=16, tickfont_size=14, gridcolor='#333333', showgrid=True),
            yaxis=dict(title_font_size=16, tickfont_size=14, gridcolor='#333333', showgrid=True),
            hovermode='x unified', margin=dict(b=130))
        return fig
    
    return {'amd_plot': create_fig(build_data('amd'), 'AMD'),
            'nvidia_plot': create_fig(build_data('nvidia'), 'NVIDIA')}