import os
import gradio as gr
import numpy as np
import matplotlib.pyplot as plt
import random
from typing import List
from rcwa import Material, Layer, LayerStack, Source, Solver
from smolagents import tool, CodeAgent, InferenceClientModel, stream_to_gradio

# --- Constants ---
start_wl = 0.32
stop_wl = 0.80
step_wl = 0.01
wavelengths = np.arange(start_wl, stop_wl + step_wl, step_wl)
materials = ['Si', 'Si3N4', 'SiO2', 'AlN']

@tool
def simulate_spectrum_100nm(layer_order: List[str]) -> List[float]:
    """
    Simulates the optical transmission spectrum for a given sequence of material layers at 100nm thickness.

    Args:
        layer_order (List[str]): A list of material names (e.g., ["Si", "SiO2", "AlN"]) representing the order of layers in the optical stack.

    Returns:
        List[float]: The transmission spectrum across a predefined wavelength range.
    """
    source = Source(wavelength=start_wl)
    reflection_layer = Layer(n=1.0)
    transmission_layer = Layer(material=Material("Si"))
    try:
        layers = [Layer(material=Material(m), thickness=0.1) for m in layer_order]
        stack = LayerStack(*layers, incident_layer=reflection_layer, transmission_layer=transmission_layer)
        solver = Solver(stack, source, (1, 1))
        result = solver.solve(wavelength=wavelengths)
        return np.array(result['TTot']).tolist()
    except Exception as e:
        return []

@tool
def cosine_similarity(vec1: List[float], vec2: List[float]) -> float:
    """
    Computes the cosine similarity between two vectors.

    Args:
        vec1 (List[float]): The first vector.
        vec2 (List[float]): The second vector.

    Returns:
        float: A similarity score between -1 and 1.
    """
    a, b = np.array(vec1), np.array(vec2)
    return float(np.dot(a, b) / (np.linalg.norm(a) * np.linalg.norm(b)))


# --- Target Spectrum Generator ---
def get_target_spectrum(layer_order, thickness=0.1):
    source = Source(wavelength=start_wl)
    reflection_layer = Layer(n=1.0)
    transmission_layer = Layer(material=Material("Si"))
    try:
        layers = [Layer(material=Material(m), thickness=thickness) for m in layer_order]
        stack = LayerStack(*layers, incident_layer=reflection_layer, transmission_layer=transmission_layer)
        solver = Solver(stack, source, (1, 1))
        result = solver.solve(wavelength=wavelengths)
        return np.array(result['TTot']).tolist()
    except Exception:
        return None

# --- Model Setup ---
from smolagents import LiteLLMModel
openai_key = os.getenv("OPENAI_API_KEY")
model = LiteLLMModel(model_id="openai/gpt-4.1-mini", temperature=0, api_key=openai_key)

# --- Agent Setup ---
agent_10nm_simulator = CodeAgent(
    tools=[simulate_spectrum_10nm],
    model=model,
    stream_outputs=True,
    name="agent_10nm_simulator",
    description="You are an AI agent that uses tools to simulate optical spectra for materials with thickness 10nm. You must provide the simulated response back. Do not provide any other information. "
)

agent_10nm_simulator.prompt_templates['managed_agent'] = {
    "task": """You're an assistant agent named '{{name}}'.
You have been given this task:
---
{{task}}
---
Just return the result of your tool call. Do not add explanations or formatting.
Call a tool immediately and use `final_answer(...)` to return the result.
""",
    "report": """{{final_answer}}"""  # Minimal required key
}



agent_100nm_simulator = CodeAgent(
    tools=[simulate_spectrum_100nm],
    model=model,
    stream_outputs=True,
    name="agent_100nm_simulator",
    description="You are an AI agent that uses tools to simulate optical spectra for materials with thickness 100nm. You must provide the simulated response back. Do not provide any other information."
)


agent_100nm_simulator.prompt_templates['managed_agent'] = {
    "task": """You're an assistant agent named '{{name}}'.
You have been given this task:
---
{{task}}
---
Just return the result of your tool call. Do not add explanations or formatting.
Call a tool immediately and use `final_answer(...)` to return the result.
""",
    "report": """{{final_answer}}"""  # Minimal required key
}


coordinator = CodeAgent(
    tools=[cosine_similarity],
    managed_agents=[agent_10nm_simulator, agent_100nm_simulator],
    model=model,
    stream_outputs=True,
    additional_authorized_imports = ["numpy"]
)

# --- Gradio UI ---
with gr.Blocks() as demo:
    gr.Markdown("""
    # 🧠 Multi-Agent Thin Film Stack Optimizer
    This demo simulates an AI agent coordinating spectrum simulations at 10nm and 100nm thickness to match a target.
    """)
    gr.Markdown("""
    CodeAgent | openai/gpt-4.1-mini  
├── ✅ Authorized imports: ['numpy']  
├── 🛠️ Tools:  
│   ┏━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┓  
│   ┃ Name              ┃ Description                                ┃ Arguments                                  ┃  
│   ┡━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┩  
│   │ cosine_similarity │ Computes the cosine similarity between two │ vec1 (`array`): The first vector.          │  
│   │                   │ vectors.                                   │ vec2 (`array`): The second vector.         │  
│   └───────────────────┴────────────────────────────────────────────┴────────────────────────────────────────────┘  
└── 🤖 Managed agents:  
    ├── agent_10nm_simulator | CodeAgent | openai/gpt-4.1-mini  
    │   ├── ✅ Authorized imports: []  
    │   ├── 📝 Description: Simulates optical spectra for 10nm thickness.  
    │   └── 🛠️ Tools:  
    │       ┏━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┓  
    │       ┃ Name                   ┃ Description                         ┃ Arguments                            ┃  
    │       ┡━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┩  
    │       │ simulate_spectrum_10nm │ Simulates spectrum for 10nm layers. │ layer_order (`array`): List of materials │  
    │       └────────────────────────┴─────────────────────────────────────┴──────────────────────────────────────┘  
    └── agent_100nm_simulator | CodeAgent | openai/gpt-4.1-mini  
        ├── ✅ Authorized imports: []  
        ├── 📝 Description: Simulates optical spectra for 100nm thickness.  
        └── 🛠️ Tools:  
            ┏━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┓  
            ┃ Name                    ┃ Description                         ┃ Arguments                           ┃  
            ┡━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┩  
            │ simulate_spectrum_100nm │ Simulates spectrum for 100nm layers.│ layer_order (`array`): List of materials │  
            └─────────────────────────┴─────────────────────────────────────┴─────────────────────────────────────┘  
    """)
    run_btn = gr.Button("🔁 Run Agent on Random Stack")
    true_order = gr.Textbox(label="True Material Order")
    prompt_box = gr.Textbox(label="Agent Prompt")
    chatbot = gr.Chatbot(label="Agent Reasoning Stream")

    def run_agent_streaming():
        true_order_val = random.sample(materials, 4)
        target_val = get_target_spectrum(true_order_val)
        true_order_display = ", ".join(true_order_val)

        if target_val is None:
            yield gr.update(value="Simulation failed"), gr.update(), gr.update()
            return

        prompt = f"""
        You are the Coordinator Agent. Your objective is to identify a 4-layer material stack **order** and **thickness** that reproduces a given target optical transmission spectrum.
        
        Constraints:
        - Materials: [Si, Si3N4, SiO2, AlN] (use each exactly once)
        - Two fixed thickness options for all layers: 10nm and 100nm
        
        You have access to the following:
        - agent_10nm_simulator: An agent that simulates a spectrum for a given material order with **10nm** layer thickness
        - agent_100nm_simulator: An agent that simulates a spectrum for a given material order with **100nm** layer thickness
        - cosine_similarity: Compares a predicted spectrum to the target spectrum
        
        Your task:
        1. Choose candidate layer orders and thickness options
        2. Call the appropriate agent to simulate the spectrum
        3. Use cosine_similarity to compare with the target
        4. Stop when similarity exceeds 0.999
        5. Report the matching order, thickness, and number of attempts
        
        Begin.
        
        Target spectrum: {target_val}
        """
        chat_history = []
        yield gr.update(value=true_order_display), gr.update(value=prompt), gr.update(value=[])

        for msg in stream_to_gradio(coordinator, task=prompt):
            if isinstance(msg, gr.ChatMessage):
                chat_history.append(("", msg.content))
            elif isinstance(msg, str):
                if chat_history:
                    chat_history[-1] = ("", msg)
                else:
                    chat_history.append(("", msg))
            yield gr.update(), gr.update(), gr.update(value=chat_history)

    run_btn.click(fn=run_agent_streaming, inputs=[], outputs=[true_order, prompt_box, chatbot])

    demo.launch(debug=True)