"""
AdaptiveDNA API — Hugging Face Space backend
FastAPI implementation of all AdaptiveDNA GPT tools.

Deploy steps:
1. Create a new HF Space (SDK: Docker or Gradio)
2. Upload this file as app.py
3. Upload requirements.txt alongside it
4. The Space URL becomes your OpenAPI server URL in schema.yaml
"""

from __future__ import annotations

import re
from typing import Literal

from fastapi import FastAPI, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel, Field

app = FastAPI(
    title="AdaptiveDNA API",
    description="CRISPR guide RNA design and DNA sequence analysis for the AdaptiveDNA GPT.",
    version="1.0.0",
)

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_methods=["*"],
    allow_headers=["*"],
)

# ── IUPAC ambiguity codes for PAM matching ────────────────────────────────────

IUPAC: dict[str, str] = {
    "N": "[ATGC]", "R": "[AG]",  "Y": "[CT]",  "S": "[GC]",
    "W": "[AT]",   "K": "[GT]",  "M": "[AC]",  "B": "[CGT]",
    "D": "[AGT]",  "H": "[ACT]", "V": "[ACG]",
}

RESTRICTION_ENZYMES: dict[str, str] = {
    "EcoRI":   "GAATTC", "BamHI":   "GGATCC", "HindIII": "AAGCTT",
    "NcoI":    "CCATGG", "XhoI":    "CTCGAG",  "NheI":    "GCTAGC",
    "SalI":    "GTCGAC", "XbaI":    "TCTAGA",  "SmaI":    "CCCGGG",
    "KpnI":    "GGTACC", "SacI":    "GAGCTC",  "ApaI":    "GGGCCC",
}

NUCLEASE_PRESETS: dict[str, dict] = {
    "SpCas9":      {"pam": "NGG",      "guide_length": 20, "description": "Most common — Streptococcus pyogenes Cas9"},
    "SaCas9":      {"pam": "NNGRRT",   "guide_length": 21, "description": "Smaller size — Staphylococcus aureus Cas9, AAV-compatible"},
    "Cpf1/Cas12a": {"pam": "TTTV",     "guide_length": 24, "description": "5′ PAM, staggered cuts — Cpf1/Cas12a, good for AT-rich regions"},
    "CjCas9":      {"pam": "NNNNRYAC", "guide_length": 22, "description": "Ultra-compact — Campylobacter jejuni Cas9"},
}

CROP_GENE_TARGETS: dict[str, list[dict]] = {
    "Rice": [
        {"gene": "OsBADH2",  "trait": "Aroma",           "description": "Knockout produces 2-AP aromatic compounds"},
        {"gene": "OsSPL14",  "trait": "Yield",            "description": "Regulation of panicle branching"},
        {"gene": "OsGW5",    "trait": "Grain width",      "description": "Controls grain width and weight"},
        {"gene": "OsDREB1",  "trait": "Drought",          "description": "Dehydration-responsive element binding"},
        {"gene": "Xa13",     "trait": "Disease resist",   "description": "Bacterial blight susceptibility gene"},
    ],
    "Maize": [
        {"gene": "ZmLG1",    "trait": "Leaf angle",       "description": "Liguleless1 — canopy architecture"},
        {"gene": "ZmKW6",    "trait": "Kernel weight",    "description": "Controls kernel width"},
        {"gene": "ZmWAKL",   "trait": "Blight resist",   "description": "Wall-associated kinase-like"},
        {"gene": "ZmDREB2A", "trait": "Heat stress",      "description": "Drought/heat tolerance transcription factor"},
    ],
    "Wheat": [
        {"gene": "TaGW2",    "trait": "Grain size",       "description": "RING-type E3 ubiquitin ligase"},
        {"gene": "TaDep1",   "trait": "Dense ears",       "description": "Dense and erect panicle orthologue"},
        {"gene": "TaMLO",    "trait": "Powdery mildew",   "description": "Mlo orthologue — disease resistance"},
        {"gene": "TaGASR7",  "trait": "Grain length",     "description": "Gibberellin-regulated"},
    ],
    "Tomato": [
        {"gene": "SlCLV3",   "trait": "Fruit size",       "description": "Clavata3 — meristem identity"},
        {"gene": "SlWUS",    "trait": "Fruit number",     "description": "Wuschel — locule number"},
        {"gene": "SlPDS",    "trait": "Carotenoids",      "description": "Phytoene desaturase — lycopene"},
        {"gene": "SlMlo1",   "trait": "PM resistance",    "description": "Powdery mildew resistance"},
    ],
    "Soybean": [
        {"gene": "GmFT2a",   "trait": "Flowering",        "description": "Florigen — photoperiod adaptation"},
        {"gene": "GmFAD2",   "trait": "Fatty acids",      "description": "Fatty acid desaturase — oleic acid"},
        {"gene": "GmPDS11",  "trait": "Pigmentation",     "description": "Phytoene desaturase visual marker"},
    ],
}

# ── Positional scoring weights (Doench 2016 / Rule Set 2 simplified) ─────────

POSITIONAL_WEIGHTS = [
    (1,  "G",  0.12), (2,  "A", -0.07), (3,  "C",  0.05), (4,  "T", -0.04),
    (12, "C",  0.09), (12, "T", -0.11), (13, "G",  0.08), (16, "G",  0.07),
    (17, "A", -0.06), (19, "G",  0.10), (20, "G",  0.08),
]

# ── Sequence utilities ────────────────────────────────────────────────────────

COMPLEMENT = str.maketrans("ATGCatgcNn", "TACGtacgNn")


def clean_sequence(raw: str) -> str:
    return re.sub(r"[^ATGCNatgcn]", "", raw).upper()


def reverse_complement(seq: str) -> str:
    return seq.translate(COMPLEMENT)[::-1]


def gc_content(seq: str) -> float:
    if not seq:
        return 0.0
    gc = sum(1 for b in seq.upper() if b in "GC")
    return round((gc / len(seq)) * 1000) / 10


def melting_temp(seq: str) -> float:
    s = seq.upper()
    at = sum(1 for b in s if b in "AT")
    gc = sum(1 for b in s if b in "GC")
    return float(2 * at + 4 * gc)


def pam_to_regex(pam: str) -> re.Pattern:
    pattern = "".join(IUPAC.get(c, c) for c in pam.upper())
    return re.compile(pattern)


def score_guide(guide: str) -> float:
    g = guide.upper()
    score = 0.5
    gc = gc_content(g) / 100
    if 0.4 <= gc <= 0.7:
        score += 0.15
    elif gc < 0.3 or gc > 0.8:
        score -= 0.25
    if re.search(r"TTTT", g):
        score -= 0.20
    if re.search(r"AAAA|CCCC|GGGG", g):
        score -= 0.10
    for pos, base, weight in POSITIONAL_WEIGHTS:
        if pos - 1 < len(g) and g[pos - 1] == base:
            score += weight
    seed = g[11:]
    seed_gc = gc_content(seed) / 100
    if 0.35 <= seed_gc <= 0.65:
        score += 0.08
    if g and g[0] == "G":
        score += 0.04
    return max(0.0, min(1.0, score))


def estimate_off_targets(guide: str) -> int:
    seed = guide[-12:].upper()
    gc_seed = gc_content(seed) / 100
    risk = 0
    if gc_seed > 0.7:
        risk += 3
    elif gc_seed > 0.55:
        risk += 1
    if re.search(r"(.{3,})\1", guide):
        risk += 2
    if re.search(r"^.{0,3}GGG|GGG.{0,3}$", guide):
        risk += 1
    return max(0, risk)


def complexity(score: float, off_targets: int) -> str:
    if score >= 0.7 and off_targets == 0:
        return "Low"
    if score >= 0.5 and off_targets <= 2:
        return "Medium"
    return "High"


def design_guides(sequence: str, pam: str, guide_len: int, top_n: int) -> list[dict]:
    seq = clean_sequence(sequence)
    pam_rx = pam_to_regex(pam)
    results: list[dict] = []

    for strand, s in [("+", seq), ("-", reverse_complement(seq))]:
        for m in pam_rx.finditer(s):
            pam_start = m.start()
            if pam_start < guide_len:
                continue
            guide = s[pam_start - guide_len: pam_start]
            sc = score_guide(guide)
            gc = gc_content(guide)
            off = estimate_off_targets(guide)
            pos = pam_start - guide_len + 1 if strand == "+" else len(seq) - pam_start + 1
            results.append({
                "sequence":    guide,
                "pam":         m.group(),
                "position":    pos,
                "strand":      strand,
                "gc_content":  gc,
                "score":       round(sc, 2),
                "off_targets": off,
                "tm_celsius":  melting_temp(guide),
                "complexity":  complexity(sc, off),
            })

    results.sort(key=lambda r: r["score"], reverse=True)
    return results[:top_n]


# ── Request / response models ─────────────────────────────────────────────────

class SequenceRequest(BaseModel):
    sequence: str = Field(..., description="DNA sequence (ATGCN characters)")


class GuideRNARequest(BaseModel):
    sequence: str = Field(..., description="Target DNA sequence (min 23 bp)")
    nuclease: Literal["SpCas9", "SaCas9", "Cpf1/Cas12a", "CjCas9"] = Field("SpCas9")
    top_n: int = Field(5, ge=1, le=10)


class CropRequest(BaseModel):
    crop: Literal["Rice", "Maize", "Wheat", "Tomato", "Soybean"]


# ── Endpoints ─────────────────────────────────────────────────────────────────

@app.post("/analyze_sequence")
def analyze_sequence(body: SequenceRequest) -> dict:
    seq = clean_sequence(body.sequence)
    if not seq:
        raise HTTPException(422, detail={"error": "Sequence is empty after cleaning"})
    s = seq.upper()
    return {
        "length":     len(seq),
        "gc_content": gc_content(seq),
        "at_content": round(((s.count("A") + s.count("T")) / len(seq)) * 1000) / 10,
        "tm_celsius": melting_temp(seq),
        "a_count":    s.count("A"),
        "t_count":    s.count("T"),
        "g_count":    s.count("G"),
        "c_count":    s.count("C"),
        "n_count":    s.count("N"),
    }


@app.post("/design_guide_rna")
def design_guide_rna(body: GuideRNARequest) -> dict:
    seq = clean_sequence(body.sequence)
    if len(seq) < 23:
        raise HTTPException(422, detail={"error": f"Sequence too short ({len(seq)} bp); minimum 23 bp required"})
    preset = NUCLEASE_PRESETS[body.nuclease]
    guides = design_guides(seq, preset["pam"], preset["guide_length"], body.top_n)
    return {"count": len(guides), "guides": guides}


@app.post("/find_crop_targets")
def find_crop_targets(body: CropRequest) -> dict:
    targets = CROP_GENE_TARGETS.get(body.crop)
    if not targets:
        raise HTTPException(422, detail={"error": f"Unknown crop: {body.crop}"})
    return {"crop": body.crop, "targets": targets}


@app.post("/calculate_gc")
def calculate_gc(body: SequenceRequest) -> dict:
    seq = clean_sequence(body.sequence)
    preview = seq[:20] + ("…" if len(seq) > 20 else "")
    return {"sequence_preview": preview, "gc_percent": gc_content(seq)}


@app.get("/nuclease_presets")
def get_nuclease_presets() -> dict:
    presets = [
        {"name": name, "pam": v["pam"], "guide_length": v["guide_length"], "description": v["description"]}
        for name, v in NUCLEASE_PRESETS.items()
    ]
    return {"presets": presets}


@app.get("/crops")
def list_crops() -> dict:
    return {"crops": list(CROP_GENE_TARGETS.keys())}


@app.post("/restriction_sites")
def find_restriction_sites(body: SequenceRequest) -> dict:
    seq = clean_sequence(body.sequence)
    sites = []
    for enzyme, site in RESTRICTION_ENZYMES.items():
        positions = []
        idx = seq.find(site)
        while idx != -1:
            positions.append(idx + 1)
            idx = seq.find(site, idx + 1)
        if positions:
            sites.append({"enzyme": enzyme, "site": site, "positions": positions})
    return {"sites": sites}


@app.get("/health")
def health() -> dict:
    return {"status": "ok", "service": "AdaptiveDNA API"}