app/logic/calc_b.py

import numpy as np
from typing import Dict, List
import google.generativeai as genai
import faiss
from openai import OpenAI
import json
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)

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 calculate_vq_dqdv(V: List[float], I: float, t: List[float]) -> Dict[str, List[float]]:
    V_arr = np.array(V, dtype=float)
    t_arr = np.array(t, dtype=float)
    Q_arr = I * t_arr
    dQ = np.gradient(Q_arr)
    dV = np.gradient(V_arr)
    with np.errstate(divide='ignore', invalid='ignore'):
        dQdV = np.where(dV!=0, dQ/dV, 0.0)
    return {"Q": Q_arr.tolist(), "V": V_arr.tolist(), "dQdV": dQdV.tolist()}

def calculate_rate_capability(
    Q_nominal: float,
    C_rates: List[float],
    t_discharge: List[float]
) -> Dict[str, List[float]]:
    # Convert to arrays
    C = np.array(C_rates, dtype=float)
    t = np.array(t_discharge, dtype=float)

    # Validation
    if C.shape != t.shape:
        raise ValueError("C_rates and t_discharge must have same length")

    # Q_Ci = C_i * Q_nominal * t_i
    Q_ci = C * Q_nominal * t
    return {"C_rates": C.tolist(), "Q_Ci": Q_ci.tolist()}


def calculate_cccv_time(
    Q_nominal: float,
    I_lim: float,
    alpha: float,
    I_end: float,
    tau: float
) -> float:
    """
    t_charge = t_CC + t_CV
             = (alpha * Q_nominal)/I_lim + tau * ln(I_lim/I_end)
    """
    # CC time:
    t_CC = (alpha * Q_nominal) / I_lim
    # CV time:
    t_CV = tau * np.log(I_lim / I_end)
    return float(t_CC + t_CV)

def calculate_diffusion(
    L: float,
    tau_pulse: float,
    delta_E_tau: List[float],
    delta_E_s: List[float]
) -> Dict[str, List[float]]:
    """
    D_i = (π/4) * (L^2 / τ_pulse) * (ΔEτ_i / ΔEs_i)^2
    """
    Δτ = np.array(delta_E_tau, dtype=float)
    Δs = np.array(delta_E_s, dtype=float)
    # avoid divide‑by‑zero
    ratio2 = np.where(Δs != 0, (Δτ / Δs)**2, 0.0)
    D = (np.pi / 4) * (L**2 / tau_pulse) * ratio2
    return {"D": D.tolist()}

def calculate_all_b(cathode_name: str, input_data: Dict) -> Dict:
    # 1) V–Q & dQ/dV
    vq = calculate_vq_dqdv(
        V=input_data["V"],
        I=input_data["I"],
        t=input_data["t"]
    )

    # 2) Rate capability
    rate = calculate_rate_capability(
        input_data["Q_nominal"],
        input_data["C_rates"],
        input_data["t_discharge"]
    )

    # 3) CC–CV charge time
    t_charge = calculate_cccv_time(
        Q_nominal = input_data["Q_nominal_mAh"],
        I_lim     = input_data["I_lim"],
        alpha     = input_data["alpha"],
        I_end     = input_data["I_end"],
        tau       = input_data["tau"]
    )

    diffusion = calculate_diffusion(
        L           = input_data["L"],
        tau_pulse   = input_data["tau_pulse"],
        delta_E_tau = input_data["delta_E_tau"],
        delta_E_s   = input_data["delta_E_s"]
    )

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

    faiss_results = query_faiss_index(query_text, top_k=5)

    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.
        And 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 of the calculations.
        Lastly shortly analyze the performance of the full cell battery.

        1. **V–Q curve**: {vq['Q']} vs {vq['V']}
        2. **dQ/dV curve**: {vq['dQdV']}
        3. **Rate capability test**:
            - C-rates: {rate['C_rates']}
            - Discharge capacities: {rate['Q_Ci']}
        4. **CC–CV charging time**: {t_charge:.2f} seconds
        5. **Diffusion coefficients** (D): {diffusion['D']}

        Please comment on:
        - Whether the electrode is rate-limited
        - Diffusion characteristics
        - Fast-charging behavior
        - Any obvious limitations or strengths

        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 the information from this section to explain the results:
        {json.dumps(faiss_results, indent=2)}
            """

    try:
        gemini_response = model.generate_content(prompt)
    
        if gemini_response.candidates:
            parts = gemini_response.candidates[0].content.parts
            gemini_text = " ".join(
                p.text for p in parts if hasattr(p, "text") and p.text
            )
        else:
            gemini_text = "Gemini returned no candidates."
    
    except Exception as e:
        gemini_text = f"Gemini analysis failed: {str(e)}"

    return {
      "vq_curve":        vq,
      "rate_capability": rate,
      "cccv_time_s":     t_charge,
      "diffusion_D":     diffusion["D"],
      "Gemini_Explanation": gemini_text
    }