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 Definitions ---
@tool
def simulate_spectrum(layer_order: List[str]) -> List[float]:
    """
    Simulates the optical transmission spectrum for a given sequence of material layers.

    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)))

# --- Agent Setup ---
from smolagents import LiteLLMModel

# You can replace this with any key (e.g. os.getenv("OPENAI_API_KEY"))
openai_key = os.getenv("OPENAI_API_KEY")  

model = LiteLLMModel(
    model_id="openai/gpt-4.1-mini",
    temperature=0,
    api_key=openai_key
)
def build_agent():
    return CodeAgent(tools=[simulate_spectrum, cosine_similarity], model=model, stream_outputs=True)

# --- Utils ---
def get_target_spectrum(layer_order):
    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:
        return None

# --- Gradio UI ---
with gr.Blocks() as demo:
    gr.Markdown("""
    # 🧠 Optical Thin-Film Stack AI Agent with Multiple Tool Calls

    This interactive demo showcases an **AI agent that coordinates multiple tool calls** to solve an inverse optics problem.  
    Rather than relying on a single inference, the agent repeatedly invokes a **physics-based simulator (RCWA)** and a **spectrum comparison function** to iteratively search for the correct material ordering in a thin-film stack.  
    It dynamically explores options, evaluates feedback, and stops only when a verifiably accurate match is found.

    ### 🤖 Objective: Reverse-engineer a 4-layer thin-film stack  
    Given only a target transmission spectrum, the agent must find the **correct material order** that would reproduce it.
    **Constraints:**  
    - Materials: `Si`, `Si₃N₄`, `SiO₂`, `AlN` (used once each)  
    - Terminate when `cosine_similarity > 0.999999`

    ## 🔍 What's Happening Under the Hood

    1. A **random 4-layer stack** is generated from a predefined material set (`Si`, `Si₃N₄`, `SiO₂`, `AlN`), with each layer set to 100nm.
    2. Its **optical transmission spectrum** is computed using **RCWA (Rigorous Coupled-Wave Analysis)** — a high-fidelity physics simulator.
    3. This target spectrum is handed to an AI **CodeAgent**, powered by GPT-4.1-mini, along with access to callable tools.
    4. The agent dynamically explores candidate layer permutations by invoking `simulate_spectrum(...)` to generate spectra, and `cosine_similarity(...)` to compare them to the target.
    5. It loops over permutations, choosing what to simulate next based on feedback — and stops automatically when similarity exceeds a threshold (e.g., `> 0.999999`).
    6. The final output includes the matched material order, a reasoning trace, and optionally a spectrum comparison plot.    
    """)
    
    gr.Markdown("### 📊 Transmission Spectra of Layer Stack Designs for all 24 permutations of material order")
    gr.Image(value="121_resized.png", interactive=False)
    gr.Markdown("""

    ### 🛠️ Tools available to the Agent: 

    | Tool               | Description                                | Arguments                              |
    |--------------------|--------------------------------------------|-----------------------------------------|
    | `simulate_spectrum`| Simulates optical spectrum                 | `layer_order`: List of materials        |
    | `cosine_similarity`| Compares predicted vs target spectra       | `vec1`, `vec2`: Lists of floats         |
    | `final_answer`     | Emits final result with justification      | `answer`: any                           |

    > 🧪 This is not a single tool call to the Physics Solver (RCWA) — the agent uses simulation as feedback to drive tool selection and termination.
    """)


    gr.Markdown("# 🧪 Optical Spectrum Matching Agent (Streaming UI Style)")
    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():
        agent = build_agent()  # ✅ fresh agent per run
        true_order_val = random.sample(materials, 4)
        target_val = get_target_spectrum(true_order_val)
        agent_summary = agent.visualize()
        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 an AI agent that uses tools to simulate optical spectra and compare them to a target spectrum.
You must find a 4-layer permutation from this list: [Si, Si3N4, SiO2, AlN]. Use each material only once.
Use simulate_spectrum(order, thicknesses) to simulate and cosine_similarity(predicted, target) to compare.
Stop when similarity > 0.999999. Report number of permutations tried.
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(agent, 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()