# strategy_llm.py
import pandas as pd
import ollama
import json
import os

# === CONFIG ===
EXCEL_FILE = "F12022ResultsMcLaren.xlsx"
MODEL_NAME = "qwen2.5:3b"
MAX_RACES_TO_PROCESS = 5   # for now first 5; 

# === SYSTEM PROMPT: teach the model real F1 strategy context ===
SYSTEM_PROMPT = """
You are an F1 race strategy engineer AI with expert knowledge of Formula One tyre compounds and pit strategies.

TYRE & STRATEGY KNOWLEDGE:
- Tyre compounds range roughly from C1 (hardest) to C5 (softest).
- Softer tyres provide more grip and faster lap times but degrade faster.
- Harder tyres are slower but more durable and better for long stints.
- Dry races typically involve 1–3 pit stops depending on tyre degradation, track abrasiveness, and fuel load.
- Hot, windy, or abrasive tracks increase tyre wear; cooler tracks reduce it.
- Finishing position: lower number = better (P1 is a win, P20 is bad).
- Missing pit stop data usually means the driver did not stop at that stage.

YOUR GOAL:
Given race data (starting tyre, pit stop laps and compounds, weather, finishing position, and notes),
you must judge how good the strategy was and suggest a realistic change that could improve the result.

You MUST think like a real race engineer and respect physical & tactical constraints.
"""

# === Build a prompt for each race ===
def turn_race_into_prompt(row):
    race = row.get("race", "Unknown Race")
    driver = row.get("driver", "Unknown Driver")
    starting_tyre = row.get("starting tyre", "N/A")

    pit1_compound = row.get("pit 1 tyre", "None")
    pit1_lap = row.get("pit 1 lap", "N/A")

    pit2_compound = row.get("pit 2 tyre", "None")
    pit2_lap = row.get("pit 2 lap", "N/A")

    pit3_compound = row.get("pit 3 tyre", "None")
    pit3_lap = row.get("pit 3 lap", "N/A")

    weather = row.get("weather conditions", "N/A")
    result = row.get("finishing position", "N/A")
    notes = row.get("strategy note", "")

    prompt = f"""
You are analyzing a single race for McLaren in the 2022 Formula 1 season.

### RACE DATA
Race: {race}
Driver: {driver}
Starting Tyre: {starting_tyre}

Pit Stops:
- Pit 1: {pit1_compound} at lap {pit1_lap}
- Pit 2: {pit2_compound} at lap {pit2_lap}
- Pit 3: {pit3_compound} at lap {pit3_lap}

Weather: {weather}
Finishing Position: {result}
Strategy Notes: {notes}

### TASK
Act as a professional F1 race strategy engineer. Evaluate this strategy and answer in STRICT JSON ONLY with EXACTLY this structure:

{{
  "overall_rating": "<good|average|poor>",
  "what_went_well": ["bullet point list of positive aspects"],
  "what_went_wrong": ["bullet point list of strategy flaws or risks"],
  "optimal_strategy_change": {{
      "pit_lap": <integer lap number or null>,
      "compound": "<recommended compound or tyre type or null>",
      "expected_gain": "<short description of expected improvement, e.g. 'better tyre life', 'lower degradation', 'track position'>",
      "reasoning": "short explanation linking tyre choice, laps, weather, and finishing position"
  }},
  "confidence": "<high|medium|low>"
}}

RULES:
- Respond with VALID JSON ONLY. No extra text, no comments, no markdown.
- If the existing strategy was already close to optimal, set "overall_rating" to "good" and keep changes minimal.
- If you can't infer a realistic change, set pit_lap and compound to null and explain why in reasoning.
"""
    return prompt

# === Run prompt through Qwen model ===
def call_model(prompt_text: str) -> str:
    response = ollama.chat(
        model=MODEL_NAME,
        messages=[
            {"role": "system", "content": SYSTEM_PROMPT},
            {"role": "user", "content": prompt_text}
        ]
    )
    return response["message"]["content"]

# === Main Program ===
def main():
    if not os.path.exists(EXCEL_FILE):
        print(f"Error: {EXCEL_FILE} not found in this folder.")
        return

    # Load Excel
    df = pd.read_excel(EXCEL_FILE)
    print("Columns found:", list(df.columns))
    print("Rows loaded:", len(df))

    # Normalize column names to lowercase + strip spaces
    df.columns = [c.lower().strip() for c in df.columns]

    # Optionally limit number of races while testing
    df_to_use = df.head(MAX_RACES_TO_PROCESS)

    results = []

    for index, row in df_to_use.iterrows():
        print(f"\n=== Processing Race {index+1}: {row.get('race', 'Unknown')} ===")

        prompt = turn_race_into_prompt(row)
        ai_raw_output = call_model(prompt)

        print("\nAI Output (raw JSON):")
        print(ai_raw_output)

        # We will just store the raw JSON string for now;
        # later we can parse and validate it.
        results.append({
            "race_index": index + 1,
            "race": row.get("race", ""),
            "driver": row.get("driver", ""),
            "model_output_json": ai_raw_output
        })

    # Save all results to a file
    out_file = "strategy_outputs.json"
    with open(out_file, "w", encoding="utf-8") as f:
        json.dump(results, f, indent=2, ensure_ascii=False)

    print(f"\n🎉 Done! Processed {len(df_to_use)} races. All strategy results saved to {out_file}")

if __name__ == "__main__":
    main()
