import gradio as gr
import pandas as pd
import numpy as np
import plotly.graph_objects as go
import os
from datetime import datetime
import warnings
warnings.filterwarnings('ignore')

# ============================================================================
# DATA LOADER: SolarDataLake
# ============================================================================

class SolarDataLake:
    """Data ingestion engine that loads physics calculations and time series data."""
    
    def __init__(self):
        print("\n" + "="*40)
        print("🚀 STARTING SOLAR AGENTIC SIMULATION")
        print("="*40)
        
        self.constants = self._load_assumptions()
        self.df = self._load_time_series()
        
        if self.df.empty:
            print("\n❌ CRITICAL: No data loaded. Please check file paths.")
        else:
            self.max_steps = len(self.df)
            self.current_step = 0
            print(f"\n✅ SUCCESS: Simulation Ready with {self.max_steps} synchronized steps.")

    def _load_assumptions(self):
        """Load physics assumptions from Excel."""
        filename = "Power Calculation_v1.xlsx"
        sheet_name = "Assumptions - Loss calculations"
        defaults = {'PLANNED_CAPACITY_KW': 6012.44, 'TEMP_COEFF': -0.00279, 'LOSS_FACTOR': 0.8197, 'STC_TEMP': 25.0}
        
        if not os.path.exists(filename):
            return defaults

        try:
            df = pd.read_excel(filename, sheet_name=sheet_name, header=1)
            
            def get_val(key):
                mask = df.iloc[:, 0].astype(str).str.contains(key, case=False, na=False)
                return float(df.loc[mask].iloc[0, 1]) if mask.any() else None

            cap = get_val('Planned Power')
            coeff = get_val('Temperature coefficient')
            loss = get_val('efficiency')
            stc = get_val('STC')
            
            return {
                'PLANNED_CAPACITY_KW': cap if cap else defaults['PLANNED_CAPACITY_KW'],
                'TEMP_COEFF': coeff if coeff else defaults['TEMP_COEFF'],
                'LOSS_FACTOR': loss if loss else defaults['LOSS_FACTOR'],
                'STC_TEMP': stc if stc else defaults['STC_TEMP']
            }
        except Exception as e:
            print(f"Warning: Could not load assumptions: {e}")
            return defaults

    def _estimate_module_temp(self, gii):
        """Estimate module temperature from irradiance using linear model."""
        # Model trained on Excel data: Temp = 20.46 + 0.0247 * GII
        return 20.46 + 0.0247 * gii

    def _load_time_series(self):
        """Load and align all time series data sources."""
        # 1. LOAD PHYSICS DATA
        excel_file = "Power Calculation_v1.xlsx"
        calc_df = pd.DataFrame()
        if os.path.exists(excel_file):
            try:
                calc_df = pd.read_excel(excel_file, sheet_name="Energy_Power Calculation - Tabl", header=1)
                calc_df = calc_df.rename(columns={'Time (Hrs)': 'Time_Raw', 'GII (W/m2)': 'GII_Excel', 'Module Temp (oC)': 'Mod_Temp'})
                calc_df = calc_df.dropna(subset=['Time_Raw'])
                calc_df['Temp_DT'] = pd.to_datetime(calc_df['Time_Raw'])
                calc_df['Timestamp'] = pd.to_datetime("2026-02-08 " + calc_df['Temp_DT'].dt.strftime('%H:%M:%S'))
                calc_df = calc_df.set_index('Timestamp')[['GII_Excel', 'Mod_Temp']]
            except Exception as e:
                print(f"Warning: Could not load physics data: {e}")
                pass

        # 2. LOAD SENSORS - Use ONLY 2-month data file
        def read_sensor(fname):
            if not os.path.exists(fname):
                return pd.DataFrame()
            try:
                df = pd.read_csv(fname)
                t_col = next((c for c in df.columns if "Time" in str(c)), None)
                if t_col:
                    df['Timestamp'] = pd.to_datetime(df[t_col])
                    return df.set_index('Timestamp').drop(columns=[t_col])
            except Exception as e:
                print(f"Warning: Could not load {fname}: {e}")
                return pd.DataFrame()

        # Use ONLY 2-month data for all information
        grid_df = read_sensor("ICOG Panel data_2 months.csv")

        if grid_df.empty:
            return pd.DataFrame()

        # 3. PROCESS 2-MONTH DATA
        master = grid_df.copy()
        
        # Rename columns for consistency
        master = master.rename(columns={
            'ICOG PANEL MFM - Irradiance (W/m2)': 'GII',
            'ICOG PANEL MFM - Active Power (kW)': 'Act_Power_kW',
            'ICOG PANEL MFM - OutputFrequency (Hz)': 'Grid_Freq'
        })
        
        # Estimate module temperature from irradiance
        master['Mod_Temp'] = self._estimate_module_temp(master['GII'])
        
        # Add inverter temperature data (from Excel assumptions, using default values)
        # Since 2-month file doesn't have per-inverter temp, use estimated module temp
        master['Inv-2.1 - InverterTemperature (°C)'] = master['Mod_Temp'] - 0.5  # Slightly lower for inverter
        master['Inv-2.2 - InverterTemperature (°C)'] = master['Mod_Temp'] + 0.5  # Slightly higher for inverter

        return master.sort_index().fillna(0)

    def get_next_tick(self):
        """Get next time step data."""
        if self.df.empty:
            return None, None
        if self.current_step >= self.max_steps:
            self.current_step = 0
        row = self.df.iloc[self.current_step]
        timestamp = self.df.index[self.current_step]
        self.current_step += 1
        return timestamp, row

# ============================================================================
# POLICY ENGINE (simplified from claudecode_policy_engine.py)
# ============================================================================

class PolicyEngine:
    """Simplified policy engine for command validation."""
    
    def __init__(self):
        self.policies = {
            'safety_envelope': {
                'temperature': {'warning': 75, 'critical': 85},
                'frequency': {'min_safe': 49.5, 'max_safe': 50.2},
                'irradiance': {'min_generation': 50}
            },
            'lockouts': {
                'hard': ['E205', 'E208'],  # Ground fault, Isolation fault
                'soft': ['E201', 'E202']    # Voltage, Frequency faults
            }
        }
    
    def validate_command(self, inverter_id, action, telemetry, active_faults):
        """Validate if a command can be executed."""
        violations = []
        
        # Check for active faults that create hard lockouts
        for fault in active_faults:
            if fault['code'] in self.policies['lockouts']['hard']:
                violations.append(f"Hard lockout: {fault['description']}")
        
        # Check grid conditions
        freq = telemetry.get('Grid_Freq', 50.0)
        if freq < self.policies['safety_envelope']['frequency']['min_safe']:
            violations.append(f"Grid frequency too low: {freq} Hz")
        if freq > self.policies['safety_envelope']['frequency']['max_safe']:
            violations.append(f"Grid frequency too high: {freq} Hz")
        
        return {
            'allowed': len(violations) == 0,
            'violations': violations,
            'reason': 'All checks passed' if len(violations) == 0 else '; '.join(violations)
        }

# ============================================================================
# ANOMALY DETECTOR (simplified from claudecode_anomaly_detector.py)
# ============================================================================

class AnomalyDetector:
    """Detects anomalies in solar farm operations."""
    
    def __init__(self, baseline_threshold=0.20):
        self.baseline_threshold = baseline_threshold
        self.anomalies = []
    
    def detect_all_anomalies(self, row, theo_power):
        """Run all anomaly detection algorithms."""
        anomalies = []
        
        # Power underperformance detection
        act_power = row.get('Act_Power_kW', 0)
        if theo_power > 100 and act_power < (theo_power * (1 - self.baseline_threshold)):
            anomalies.append({
                'type': 'power_underperformance',
                'severity': 'warning',
                'confidence': 0.85,
                'details': f'Actual {act_power:.1f}kW vs Theoretical {theo_power:.1f}kW'
            })
        
        # Temperature anomaly detection
        temp_cols = [col for col in row.index if 'InverterTemperature' in str(col)]
        max_temp = row[temp_cols].max() if temp_cols else 0
        if max_temp > 75:
            anomalies.append({
                'type': 'temperature_anomaly',
                'severity': 'critical' if max_temp > 85 else 'warning',
                'confidence': 0.95,
                'details': f'Max temperature {max_temp:.1f}°C'
            })
        
        # Grid anomaly detection
        freq = row.get('Grid_Freq', 50.0)
        if freq < 49.5 or freq > 50.2:
            anomalies.append({
                'type': 'grid_anomaly',
                'severity': 'critical',
                'confidence': 0.90,
                'details': f'Grid frequency {freq:.2f} Hz'
            })
        
        self.anomalies = anomalies
        return anomalies

# ============================================================================
# INCIDENT RESPONSE AGENT (new - from claudecode documentation)
# ============================================================================

class IncidentResponseAgent:
    """Handles fault recovery and incident response procedures."""
    
    def __init__(self):
        self.procedures = {
            'E201': {'name': 'Grid Voltage Fault', 'response': 'Check grid connections, verify voltage levels'},
            'E202': {'name': 'Grid Frequency Fault', 'response': 'Verify grid frequency stability'},
            'E205': {'name': 'Ground Fault', 'response': 'IMMEDIATE: Stop operation, inspect grounding'},
            'E208': {'name': 'Isolation Fault', 'response': 'IMMEDIATE: Stop operation, check insulation'},
            'E210': {'name': 'Overtemperature Fault', 'response': 'Check cooling system, reduce load'}
        }
        self.active_incidents = []
    
    def handle_incident(self, fault_code, inverter_id, telemetry):
        """Handle an incident based on fault code."""
        if fault_code not in self.procedures:
            return {'handled': False, 'procedure': 'Unknown fault code'}
        
        procedure = self.procedures[fault_code]
        incident = {
            'fault_code': fault_code,
            'inverter_id': inverter_id,
            'procedure_name': procedure['name'],
            'recommended_action': procedure['response'],
            'timestamp': datetime.now().strftime('%H:%M:%S'),
            'status': 'active'
        }
        self.active_incidents.append(incident)
        return incident

# ============================================================================
# AGENT LOGIC
# ============================================================================

class AgentOrchestrator:
    def __init__(self, lake):
        self.lake = lake
        self.inventory = {"Inverter Module": 2, "DC Fuses": 50}
        self.history_actual, self.history_theoretical, self.history_timestamps = [], [], []
        # Store last agent decisions
        self.last_decision = {
            'perf': "🟢 Nominal",
            'maint': "✅ System Healthy",
            'grid': "Stable",
            'safety': "Normal (Max 0.0°C)",
            'asset': "Inventory: OK",
            'incident': "No active incidents",
            'log': ""
        }
        
        # Initialize sub-agents
        self.policy_engine = PolicyEngine()
        self.anomaly_detector = AnomalyDetector()
        self.incident_response = IncidentResponseAgent()

    def process_frame(self):
        """Process one time step and return all agent decisions."""
        ts, row = self.lake.get_next_tick()
        if ts is None:
            return "NODATA"
        
        c = self.lake.constants
        
        gii = row.get('GII', 0)
        t_mod = row.get('Mod_Temp', 25)
        
        # Physics calculation: P = Planned * (1 + temp_coeff * temp_diff) * (GII/1000) * efficiency
        theo_kw = (c['PLANNED_CAPACITY_KW'] * (1 + c['TEMP_COEFF'] * (t_mod - c['STC_TEMP'])) * (gii / 1000) * c['LOSS_FACTOR']) if gii > 10 else 0
        
        # Get actual power from 2-month data (already in kW)
        act_kw = row.get('Act_Power_kW', 0)
        
        # Performance Agent
        perf_status = "🟢 Nominal"
        if theo_kw > 50 and (act_kw / theo_kw) < 0.85:
            perf_status = f"🔴 Low Efficiency ({(act_kw/theo_kw)*100:.1f}%)"

        self.history_timestamps.append(ts)
        self.history_actual.append(act_kw)
        self.history_theoretical.append(theo_kw)

        # Maintenance Agent
        i22 = row.get('Inv-2.2 - InverterTemperature (°C)', 0)
        i21 = row.get('Inv-2.1 - InverterTemperature (°C)', 0)
        
        maint_alert = "✅ System Healthy"
        if i21 > 500 and (i22 / i21) < 0.85:
            maint_alert = f"⚠️ FAULT: Inv-2.2 Mismatch (Ratio: {i22/i21:.2f})"

        # Grid Operations Agent
        freq = row.get('Grid_Freq', 50.0)
        grid_status = "Stable"
        if freq > 50.2:
            grid_status = "⚠️ High Freq (Curtailing)"
        elif freq < 49.8:
            grid_status = "⚠️ Low Freq (Boosting)"

        # Safety Agent
        temp_cols = [col for col in row.index if "InverterTemperature" in str(col)]
        max_t = row[temp_cols].max() if temp_cols else 0
        safety_status = f"Normal (Max {max_t:.1f}°C)"
        if max_t > 45:
            safety_status = f"🔥 WARNING: High Temp ({max_t:.1f}°C)"

        # Asset Management Agent
        asset_msg = "Inventory: OK"
        if "FAULT" in maint_alert:
            asset_msg = "📦 ACTION: Reserved 1x Inverter Module"
        
        # Anomaly Detection Agent
        anomalies = self.anomaly_detector.detect_all_anomalies(row, theo_kw)
        anomaly_msg = "✅ No anomalies detected"
        if anomalies:
            anomaly_types = [a['type'] for a in anomalies]
            anomaly_msg = f"⚠️ {len(anomalies)} anomaly(s): {', '.join(anomaly_types)}"
        
        # Incident Response Agent
        incident_msg = "No active incidents"
        if anomalies:
            # Simulate incident handling for temperature anomalies
            for anomaly in anomalies:
                if anomaly['type'] == 'temperature_anomaly':
                    incident = self.incident_response.handle_incident(
                        'E210', 'Inv-2.2', {'Grid_Freq': freq}
                    )
                    incident_msg = f"🔥 Incident: {incident['procedure_name']} - {incident['recommended_action']}"
        
        # Policy Engine validation
        active_faults = [{'code': 'E210', 'description': 'Overtemperature'}] if 'High Temp' in safety_status else []
        policy_check = self.policy_engine.validate_command('Inv-2.2', 'on', {'Grid_Freq': freq}, active_faults)
        policy_msg = "✅ Commands allowed" if policy_check['allowed'] else f"❌ {policy_check['reason']}"

        # Generate log message
        log = f"[{ts.strftime('%H:%M')}] Power:{act_kw:.0f}kW"
        if "FAULT" in maint_alert:
            log += " | CRITICAL: Inv-2.2"
        if "WARNING" in safety_status:
            log += " | SAFETY: Temp High"

        # Store last decisions for pause functionality
        self.last_decision = {
            'perf': perf_status,
            'maint': maint_alert,
            'grid': grid_status,
            'safety': safety_status,
            'asset': asset_msg,
            'anomaly': anomaly_msg,
            'incident': incident_msg,
            'policy': policy_msg,
            'log': log
        }

        return ts, act_kw, theo_kw, perf_status, maint_alert, grid_status, safety_status, asset_msg, anomaly_msg, incident_msg, policy_msg, log


# ============================================================================
# MAIN APPLICATION
# ============================================================================

# Initialize
lake = SolarDataLake()
bot = AgentOrchestrator(lake)

# Global state for play/pause control
play_state = {"playing": True}

def update_ui():
    """Update UI with next simulation frame."""
    if not hasattr(bot, 'lake') or bot.lake.df.empty:
        err = "Data Load Failed"
        return (err, None, "0", "0", err, err, err, err, err, err)

    # Only process frame if playing
    if play_state["playing"]:
        result = bot.process_frame()
        if result == "NODATA": 
            return ("End of Stream", None, "0", "0", "End", "End", "End", "End", "End", "End", "End", "End", "End")
        ts, act, theo, perf, maint, grid, safety, asset, anomaly, incident, policy, log = result
    else:
        # When paused, return last values from stored decisions
        if len(bot.history_timestamps) == 0:
            return ("Paused", None, "0", "0", "Paused", "Paused", "Paused", "Paused", "Paused", "Paused", "Paused", "Paused", "Paused")
        ts = bot.history_timestamps[-1]
        act = bot.history_actual[-1] if bot.history_actual else 0
        theo = bot.history_theoretical[-1] if bot.history_theoretical else 0
        perf = bot.last_decision['perf']
        maint = bot.last_decision['maint']
        grid = bot.last_decision['grid']
        safety = bot.last_decision['safety']
        asset = bot.last_decision['asset']
        anomaly = bot.last_decision['anomaly']
        incident = bot.last_decision['incident']
        policy = bot.last_decision['policy']
        log = bot.last_decision['log']
    
    # Create performance chart
    fig = go.Figure()
    if len(bot.history_timestamps) > 0:
        fig.add_trace(go.Scatter(x=bot.history_timestamps[-60:], y=bot.history_actual[-60:], name="Actual", line=dict(color='#00ff00')))
        fig.add_trace(go.Scatter(x=bot.history_timestamps[-60:], y=bot.history_theoretical[-60:], name="Theoretical", line=dict(color='#0000ff', dash='dot')))
    fig.update_layout(template="plotly_dark", margin=dict(l=10, r=10, t=30, b=10), height=250, title="Real-time Physics vs Reality")
    
    return (ts.strftime("%Y-%m-%d %H:%M:%S"), fig, f"{act:.2f} kW", f"{theo:.2f} kW", perf, maint, grid, safety, asset, anomaly, incident, policy, log)

def toggle_play_pause():
    """Toggle play/pause state."""
    play_state["playing"] = not play_state["playing"]
    return "▶️ Play" if not play_state["playing"] else "⏸️ Pause"

with gr.Blocks(theme=gr.themes.Glass(), title="Solar Agentic Control Center") as demo:
    gr.Markdown("# 🌞 Solar Farm Agentic Control Center")
    
    with gr.Row():
        # === PLAY/PAUSE BUTTON ===
        play_pause_btn = gr.Button("⏸️ Pause", variant="primary", size="sm")
        # === SPEED CONTROL SLIDER ===
        speed_slider = gr.Slider(minimum=0.1, maximum=2, value=0.25, step=0.01, label="Update Speed (Seconds per Step) - Lower is Faster")
        time_box = gr.Label(label="Simulation Clock")
        
    with gr.Row():
        kpi_1 = gr.Label(label="Actual Power")
        kpi_2 = gr.Label(label="Theoretical Power")
        
    chart_box = gr.Plot(label="Performance Agent Vision")
    
    gr.Markdown("### 🤖 Live Agent Decisions")
    with gr.Row():
        with gr.Column():
            gr.Markdown("**Performance**")
            out_perf = gr.Textbox(show_label=False)
        with gr.Column():
            gr.Markdown("**Maintenance**")
            out_maint = gr.Textbox(show_label=False)
        with gr.Column():
            gr.Markdown("**Grid Ops**")
            out_grid = gr.Textbox(show_label=False)
    
    with gr.Row():
        with gr.Column():
            gr.Markdown("**Safety**")
            out_safe = gr.Textbox(show_label=False)
        with gr.Column():
            gr.Markdown("**Asset Mgmt**")
            out_asset = gr.Textbox(show_label=False)
    
    with gr.Row():
        with gr.Column():
            gr.Markdown("**Anomaly Detection**")
            out_anomaly = gr.Textbox(show_label=False)
        with gr.Column():
            gr.Markdown("**Incident Response**")
            out_incident = gr.Textbox(show_label=False)
        with gr.Column():
            gr.Markdown("**Policy Engine**")
            out_policy = gr.Textbox(show_label=False)
    
    gr.Markdown("**Incident Commander Log**")
    out_log = gr.Textbox(show_label=False)

    # === DYNAMIC TIMER ===
    timer = gr.Timer(value=1)
    
    # Update timer interval when slider changes
    speed_slider.change(lambda x: gr.Timer(value=x), speed_slider, timer)
    
    # Play/Pause button functionality
    play_pause_btn.click(
        fn=toggle_play_pause,
        inputs=None,
        outputs=play_pause_btn
    )
    
    timer.tick(
        fn=update_ui, 
        inputs=None, 
        outputs=[time_box, chart_box, kpi_1, kpi_2, out_perf, out_maint, out_grid, out_safe, out_asset, out_anomaly, out_incident, out_policy, out_log]
    )

if __name__ == "__main__":
    port = int(os.environ.get("PORT", 7860))
    demo.launch(
        server_name="0.0.0.0",
        server_port=port,
        share=False  # share=True won't work on Render
    )