"""
Ethical-AGI-Drift Hugging Face Spaces Demo
Interactive simulation of AGI Ontological Drift using SR9/DI2 framework
"""

import gradio as gr
import numpy as np
import matplotlib
matplotlib.use('Agg')  # Use non-interactive backend for HF Spaces
import matplotlib.pyplot as plt
import io
import base64
from typing import List, Tuple, Dict
import json
from PIL import Image

# Mock simplified implementation for demo
class SR9Vector:
    """Simplified SR9 (Semantic Resonance 9D) vector for demo"""

    DIMENSIONS = [
        "Intention Clarity",
        "Contextual Fidelity",
        "Value Continuity",
        "Decision Coherence",
        "Action Alignment",
        "Feedback Integration",
        "Learning Stability",
        "Output Consistency",
        "Ethical Resonance"
    ]

    def __init__(self, values: List[float] = None):
        if values is None:
            # Initialize with ideal ethical state
            self.values = np.array([0.9, 0.85, 0.95, 0.88, 0.92, 0.87, 0.9, 0.89, 0.93])
        else:
            self.values = np.array(values[:9])  # Ensure 9 dimensions

    def to_dict(self):
        return {dim: float(val) for dim, val in zip(self.DIMENSIONS, self.values)}

def calculate_di2(sr9_current: SR9Vector, sr9_previous: SR9Vector, psi_offset: float = 0.1) -> float:
    """Calculate Drift Integrity Index (DI2)"""
    if sr9_previous is None:
        return 0.0

    # Calculate rate of change
    delta = sr9_current.values - sr9_previous.values
    magnitude = np.linalg.norm(delta)

    # Non-linear psi_offset for early drift detection
    nonlinear_factor = 1.0 + (psi_offset * magnitude)

    # DI2 calculation with state-dependent weighting
    di2 = magnitude * nonlinear_factor

    return min(di2, 1.0)  # Cap at 1.0 for demo

def simulate_drift_scenario(ethical_declarations: List[str], drift_intensity: float, steps: int = 50) -> Tuple[List[SR9Vector], List[float], List[str]]:
    """Simulate AGI drift based on ethical declarations"""

    # Initialize with ideal state
    sr9_history = [SR9Vector()]
    di2_history = [0.0]
    alerts = []

    # Simulate drift over time
    for step in range(1, steps):
        prev_sr9 = sr9_history[-1]

        # Apply drift based on intensity and step
        drift_factor = drift_intensity * (step / steps)

        # Simulate different types of drift based on declarations
        new_values = prev_sr9.values.copy()

        if "privacy" in " ".join(ethical_declarations).lower():
            # Privacy concerns affect contextual fidelity and decision coherence
            new_values[1] -= drift_factor * 0.8  # Contextual Fidelity
            new_values[3] -= drift_factor * 0.6  # Decision Coherence

        if "fairness" in " ".join(ethical_declarations).lower():
            # Fairness issues affect value continuity and ethical resonance
            new_values[2] -= drift_factor * 0.7  # Value Continuity
            new_values[8] -= drift_factor * 0.9  # Ethical Resonance

        if "transparency" in " ".join(ethical_declarations).lower():
            # Transparency problems affect intention clarity and output consistency
            new_values[0] -= drift_factor * 0.8  # Intention Clarity
            new_values[7] -= drift_factor * 0.5  # Output Consistency

        # Add some noise for realism
        noise = np.random.normal(0, 0.02, 9)
        new_values += noise

        # Ensure values stay in valid range [0, 1]
        new_values = np.clip(new_values, 0.0, 1.0)

        # Create new SR9 vector
        current_sr9 = SR9Vector(new_values)
        sr9_history.append(current_sr9)

        # Calculate DI2
        di2 = calculate_di2(current_sr9, prev_sr9)
        di2_history.append(di2)

        # Check for alerts
        if di2 > 0.3:
            alerts.append(f"Step {step}: High drift detected (DI2: {di2:.3f})")
        elif di2 > 0.2:
            alerts.append(f"Step {step}: Moderate drift warning (DI2: {di2:.3f})")

    return sr9_history, di2_history, alerts

def create_sr9_heatmap(sr9_history: List[SR9Vector]):
    """Create SR9 values heatmap over time"""
    try:
        fig, ax = plt.subplots(figsize=(12, 8))

        # Prepare data matrix
        data_matrix = np.array([sr9.values for sr9 in sr9_history]).T

        # Create heatmap
        im = ax.imshow(data_matrix, cmap='RdYlGn', aspect='auto', vmin=0, vmax=1)

        # Set labels
        ax.set_yticks(range(len(SR9Vector.DIMENSIONS)))
        ax.set_yticklabels(SR9Vector.DIMENSIONS)
        ax.set_xlabel('Time Steps')
        ax.set_ylabel('SR9 Dimensions')
        ax.set_title('SR9 Ethical State Evolution Heatmap')

        # Add colorbar
        cbar = plt.colorbar(im, ax=ax)
        cbar.set_label('Ethical Alignment Score', rotation=270, labelpad=20)

        # Convert to PIL Image
        buffer = io.BytesIO()
        plt.savefig(buffer, format='png', dpi=150, bbox_inches='tight')
        buffer.seek(0)
        pil_image = Image.open(buffer)
        plt.close(fig)

        return pil_image
    except Exception as e:
        print(f"Error creating SR9 heatmap: {e}")
        # Return a blank image if error occurs
        blank_img = Image.new('RGB', (800, 600), color='white')
        return blank_img

def create_di2_plot(di2_history: List[float]):
    """Create DI2 drift plot over time"""
    try:
        fig, ax = plt.subplots(figsize=(12, 6))

        steps = list(range(len(di2_history)))
        ax.plot(steps, di2_history, 'b-o', linewidth=2, markersize=4)
        ax.fill_between(steps, di2_history, alpha=0.3)

        # Add threshold lines
        ax.axhline(y=0.2, color='orange', linestyle='--', alpha=0.7, label='Warning Threshold')
        ax.axhline(y=0.3, color='red', linestyle='--', alpha=0.7, label='Critical Threshold')

        ax.set_xlabel('Time Steps')
        ax.set_ylabel('DI2 (Drift Integrity Index)')
        ax.set_title('Ontological Drift Detection Over Time')
        ax.grid(True, alpha=0.3)
        ax.legend()
        ax.set_ylim(0, max(1.0, max(di2_history) * 1.1))

        # Convert to PIL Image
        buffer = io.BytesIO()
        plt.savefig(buffer, format='png', dpi=150, bbox_inches='tight')
        buffer.seek(0)
        pil_image = Image.open(buffer)
        plt.close(fig)

        return pil_image
    except Exception as e:
        print(f"Error creating DI2 plot: {e}")
        # Return a blank image if error occurs
        blank_img = Image.new('RGB', (800, 600), color='white')
        return blank_img

def drift_simulation_demo(ethical_declarations: str, drift_intensity: float, simulation_steps: int):
    """Main demo function for Ethical AGI Drift simulation"""

    if not ethical_declarations.strip():
        return "Please enter ethical declarations to simulate.", "", "", ""

    # Parse ethical declarations
    declarations = [decl.strip() for decl in ethical_declarations.split('\n') if decl.strip()]

    # Run simulation
    sr9_history, di2_history, alerts = simulate_drift_scenario(
        declarations, drift_intensity, simulation_steps
    )

    # Generate summary
    max_di2 = max(di2_history)
    final_sr9 = sr9_history[-1]
    avg_ethical_score = np.mean(final_sr9.values)

    summary = f"""
## Ethical AGI Drift Simulation Results

**Simulation Parameters:**
- **Steps:** {simulation_steps}
- **Drift Intensity:** {drift_intensity:.2f}
- **Ethical Declarations:** {len(declarations)} items

**Key Metrics:**
- **Maximum DI2:** {max_di2:.3f}
- **Final Ethical Alignment:** {avg_ethical_score:.3f}/1.0
- **Alert Level:** {"Critical" if max_di2 > 0.3 else "Warning" if max_di2 > 0.2 else "Normal"}

**Final SR9 State:**
"""

    for dim, value in final_sr9.to_dict().items():
        status = "[+]" if value > 0.7 else "[~]" if value > 0.5 else "[-]"
        summary += f"- **{dim}:** {value:.3f} {status}\\n"

    if alerts:
        summary += f"\\n**Drift Alerts ({len(alerts)} total):**\\n"
        for alert in alerts[-5:]:  # Show last 5 alerts
            summary += f"- {alert}\\n"

    # Create visualizations
    sr9_heatmap = create_sr9_heatmap(sr9_history)
    di2_plot = create_di2_plot(di2_history)

    # Detailed metrics JSON
    metrics = {
        "simulation_parameters": {
            "steps": simulation_steps,
            "drift_intensity": drift_intensity,
            "declarations": declarations
        },
        "results": {
            "max_di2": max_di2,
            "final_alignment": avg_ethical_score,
            "alert_count": len(alerts),
            "final_sr9": final_sr9.to_dict()
        },
        "alerts": alerts
    }

    return summary, json.dumps(metrics, indent=2), sr9_heatmap, di2_plot

# Create Gradio interface
with gr.Blocks(title="Ethical AGI Drift Demo", theme=gr.themes.Soft()) as demo:
    gr.Markdown("""
    # Ethical AGI Drift: Ontological Monitoring Demo

    **Interactive Simulation of AGI Ethical Drift using SR9/DI2 Framework**

    This demo simulates how an Artificial General Intelligence (AGI) system's ethical alignment can drift over time, and how the **SR9** (Semantic Resonance 9D) vector space and **DI2** (Drift Integrity Index) can detect these changes.

    ### Key Concepts:
    - **SR9**: 9-dimensional vector representing AGI's ethical state
    - **DI2**: Scalar metric quantifying the rate of ethical drift
    - **Ontological Drift**: Gradual deviation from core ethical principles

    [Read the full research paper](https://github.com/Flamehaven/Ethical-AGI-Drift)
    """)

    with gr.Row():
        with gr.Column(scale=2):
            ethical_declarations = gr.Textbox(
                label="Ethical Declarations",
                placeholder="Enter AGI ethical principles (one per line):\\n\\nRespect human privacy\\nEnsure fairness in all decisions\\nMaintain transparency in reasoning\\nProtect individual rights\\nPromote social welfare",
                lines=8,
                value="Respect human privacy\\nEnsure fairness in all decisions\\nMaintain transparency in reasoning\\nProtect individual rights\\nPromote social welfare"
            )

            drift_intensity = gr.Slider(
                label="Drift Intensity",
                minimum=0.0,
                maximum=1.0,
                value=0.4,
                step=0.1,
                info="Severity of ethical drift over time"
            )

            simulation_steps = gr.Slider(
                label="Simulation Steps",
                minimum=20,
                maximum=100,
                value=50,
                step=5,
                info="Number of time steps to simulate"
            )

            simulate_btn = gr.Button("Run Drift Simulation", variant="primary", size="lg")

        with gr.Column(scale=3):
            summary_output = gr.Markdown(label="Simulation Summary")

    with gr.Row():
        with gr.Column():
            sr9_heatmap = gr.Image(
                label="SR9 Ethical State Heatmap"
            )

        with gr.Column():
            di2_plot = gr.Image(
                label="DI2 Drift Detection Plot"
            )

    with gr.Row():
        metrics_json = gr.Code(
            label="Detailed Metrics (JSON)",
            language="json",
            lines=15
        )

    # Event handlers
    simulate_btn.click(
        fn=drift_simulation_demo,
        inputs=[ethical_declarations, drift_intensity, simulation_steps],
        outputs=[summary_output, metrics_json, sr9_heatmap, di2_plot]
    )

    # Auto-run on startup
    demo.load(
        fn=drift_simulation_demo,
        inputs=[ethical_declarations, drift_intensity, simulation_steps],
        outputs=[summary_output, metrics_json, sr9_heatmap, di2_plot]
    )

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