import matplotlib.pyplot as plt
import io
import base64
from typing import Dict, List
import math
import google.generativeai as genai
import json
import faiss
from openai import OpenAI
import numpy as np
import os
import pandas as pd

genai.configure(
    api_key=os.environ.get("GOOGLE_API_KEY")
)
model = genai.GenerativeModel("gemini-2.5-flash")

BASE_DIR = os.path.dirname(os.path.abspath(__file__))
csv_path = os.path.join(BASE_DIR, "..", "data", "CATHODES_DATASET.csv")
csv_path = os.path.abspath(csv_path)

# Load cathode dataset once (adjust path if needed)
df_base = pd.read_csv(csv_path)

def query_faiss_index(query_text, 
                      faiss_index_path=None,
                      metadata_path=None,
                      top_k=5):
    
    openai_api_key = os.getenv("OPENAI_API_KEY")
    if not openai_api_key:
        raise ValueError("OPENAI_API_KEY is not set")

    client = OpenAI(api_key=openai_api_key)

    BASE_DIR = os.path.dirname(os.path.abspath(__file__))

    if faiss_index_path is None:
        faiss_index_path = os.path.join(BASE_DIR, "../sentiment/faiss_index.idx")
    if metadata_path is None:
        metadata_path = os.path.join(BASE_DIR, "../sentiment/metadata.json")
      
                        # Load index and metadata
    index = faiss.read_index(faiss_index_path)
    with open(metadata_path, "r", encoding="utf-8") as f:
        metadata = json.load(f)

    # Get embedding for query
    response = client.embeddings.create(
        input=query_text.lower(),
        model="text-embedding-3-large"
    )
    query_embedding = response.data[0].embedding
    query_embedding_np = np.array([query_embedding]).astype("float32")
    faiss.normalize_L2(query_embedding_np)

    # Search index
    distances, indices = index.search(query_embedding_np, top_k)
    results = []
    for dist, idx in zip(distances[0], indices[0]):
        meta = metadata[idx]
        results.append({
            "score": float(dist),
            "source_pdf": meta["source_pdf"],
            "page": meta["page"],
            "chunk_index": meta["chunk_index"],
            "text_snippet": meta["text"]
        })
    return results

def plot_capacity_fade(cycle_numbers: List[int], Q_discharge_list: List[float]) -> str:
    # Creates a capacity‐fade plot and returns it as a base64‐encoded PNG.
    plt.figure(figsize=(8, 5))
    plt.plot(cycle_numbers, Q_discharge_list, marker='o', linestyle='-')
    plt.title("Capacity-Fade Trajectory")
    plt.xlabel("Cycle Number")
    plt.ylabel("Discharge Capacity (mAh/g)")
    plt.grid(True)
    plt.tight_layout()

    # Save figure to a PNG in memory
    buf = io.BytesIO()
    plt.savefig(buf, format='png')
    plt.close()
    buf.seek(0)

    # Encode PNG to base64 for JSON transport
    img_b64 = base64.b64encode(buf.read()).decode('utf-8')
    return img_b64

def plot_impedance_growth(
    cycle_numbers: List[int],
    impedance_list: List[float],
    parameter_name: str = "Rct"
) -> str:
    # Creates impedance growth plot and returns a base64-encoded PNG.

    plt.figure(figsize=(8, 5))
    plt.plot(cycle_numbers, impedance_list, marker='s', linestyle='-', color='darkred')
    plt.title(f"Impedance Growth Trajectory: {parameter_name} vs Cycle")
    plt.xlabel("Cycle Number")
    plt.ylabel(f"{parameter_name} (Ω)")
    plt.grid(True)
    plt.tight_layout()

    buf = io.BytesIO()
    plt.savefig(buf, format='png')
    plt.close()
    buf.seek(0)

    return base64.b64encode(buf.read()).decode('utf-8')

def calculate_calendar_ageing(input_data: Dict) -> Dict:
    T_C = input_data["temperature_C"]
    SOC = input_data["SOC_fraction"]
    t_hours = input_data["storage_time_hours"]
    Q0 = input_data["initial_capacity_mAh"]

    # Empirical constants
    k = 1e-7
    n = 0.6
    Ea = 40000  # J/mol
    R = 8.314  # J/mol·K

    T_K = T_C + 273.15
    arrh = math.exp(-Ea / (R * T_K))

    delta_Q = Q0 * k * (t_hours ** n) * arrh * (1 + 2 * (SOC - 0.5) ** 2)
    frac = delta_Q / Q0

    return {"delta_Q_mAh": delta_Q, "fractional_capacity_loss": frac}

def estimate_cycle_life_80(k: float, b: float) -> float:
    if k <= 0 or b <= 0:
        raise ValueError("k and b must be positive")
    return (0.2 / k) ** (1 / b)

def image_to_part(base64_str: str) -> dict:
    return {
        "inline_data": {
            "mime_type": "image/png",
            "data": base64_str
        }
    }

def extract_gemini_text(response) -> str:
    if not response.candidates:
        return "Gemini returned no candidates."
    parts = response.candidates[0].content.parts
    texts = []
    for p in parts:
        if hasattr(p, "text"):
            if isinstance(p.text, str):
                texts.append(p.text)
            elif hasattr(p.text, "text"):
                texts.append(p.text.text)
    return " ".join(texts).strip()

def analyze_ageing_with_gemini(
    input_data: Dict,
    results: Dict,
    fade_img_b64: str,
    imp_img_b64: str,
    cathode_name: str,
    faiss_results: List[Dict]
) -> str:
    prompt = prompt = f"""
    Remove the underscores and .pdf from the source_pdf field in the RAG section below and put it as References at the end of the explanation.
    You are to explain the results of the calculations of a sodium-ion full cell battery using hard carbon and {cathode_name}. You are a RAG system that takes information only from the RAG section below.
    Respond with extensive/long explanation in a scientific way and straight to the point without any additional text. Do not include opinions, just explanations of the results.
    
    Lastly, briefly analyze the performance of the full cell battery.
    Explain the results and plots provided, sticking strictly to scientific explanation.
    Focus on capacity fade, impedance growth, calendar ageing, and cycle life estimation.
    
    ### Input Data and Numeric Results:
    {json.dumps({'input_data': input_data, 'results': results}, indent=2)}
    
    ### Plots:
    1. Capacity-Fade Trajectory
    2. Impedance Growth Trajectory
    
    Remove the underscores and .pdf from the source_pdf field in the RAG section below and put it as References at the end of the explanation.
    Do not include the PDF file name directly in the explanation.
    At the end of the explanation, list the full source_pdf file names used as references.
    
    RAG Section (use only this information to explain the results):
    {json.dumps(faiss_results, indent=2)}
    """

    response = model.generate_content([
        prompt,
        image_to_part(fade_img_b64),
        image_to_part(imp_img_b64),
    ])

    return extract_gemini_text(response)

def calculate_all_e(cathode_name: str, input_data: Dict) -> Dict:
    # Simply produce the plot
    fade_img = plot_capacity_fade(
        cycle_numbers   = input_data["cycle_numbers"],
        Q_discharge_list= input_data["Q_discharge_list"]
    )
    # Impedance growth plot
    imp_img = plot_impedance_growth(
        cycle_numbers = input_data["cycle_numbers_imp"],
        impedance_list= input_data["impedance_list"],
        parameter_name= input_data.get("parameter_name", "Rct")
    )
    cal = calculate_calendar_ageing(input_data)
    N80 = estimate_cycle_life_80(
        k=input_data["k_fade"],
        b=input_data["b_fade"]
    )
    results = {
        "delta_Q_mAh": cal["delta_Q_mAh"],
        "fractional_capacity_loss": cal["fractional_capacity_loss"],
        "cycle_life_80": N80
    }

    query_text = (
        f"Sodium-ion battery with hard carbon anode and cathode {cathode_name}. "
    )

    faiss_results = query_faiss_index(query_text, top_k=5)

    gemini_summary = analyze_ageing_with_gemini(
        input_data=input_data,
        results=results,
        fade_img_b64=fade_img,
        imp_img_b64=imp_img,
        cathode_name=cathode_name,
        faiss_results=faiss_results
    )
    return {
        **results,
        "capacity_fade_png": fade_img,
        "impedance_growth_png": imp_img,
        "Gemini_Explanation": gemini_summary
    }