from fastapi import FastAPI
from pydantic import BaseModel
from typing import List
from nupack import Model, Strand, Tube, SetSpec, tube_analysis

app = FastAPI()

class ModelInput(BaseModel):
    material: str = "rna"
    ensemble: str = "stacking"
    celsius: float = 37
    sodium: float = 1.0
    magnesium: float = 0.0

class StrandInput(BaseModel):
    name: str
    sequence: str
    concentration: float
    unit: str = "uM"

class TubeInput(BaseModel):
    name: str = "tube1"
    max_size: int = 2

class AnalysisPayload(BaseModel):
    model: ModelInput
    strands: List[StrandInput]
    tube: TubeInput

def unit_to_molar(value: float, unit: str) -> float:
    mapping = {"M": 1, "mM": 1e-3, "uM": 1e-6, "nM": 1e-9, "pM": 1e-12}
    return value * mapping[unit]

@app.get("/health")
def health():
    return {"status": "ok"}

@app.post("/run")
def run(payload: AnalysisPayload):
    model = Model(
        material=payload.model.material,
        ensemble=payload.model.ensemble,
        celsius=payload.model.celsius,
        sodium=payload.model.sodium,
        magnesium=payload.model.magnesium,
    )

    strands = {}
    for s in payload.strands:
        obj = Strand(s.sequence, name=s.name)
        strands[obj] = unit_to_molar(s.concentration, s.unit)

    tube = Tube(
        strands=strands,
        complexes=SetSpec(max_size=payload.tube.max_size),
        name=payload.tube.name,
    )

    result = tube_analysis(tubes=[tube], model=model)

    return {
        "status": "success",
        "result": str(result)
    }