"""
RAScore micro-service for MolParetoLab (isolated from the ADMET-AI Space).

Retrosynthetic Accessibility score (Thakkar et al., Chem. Sci. 2021,
doi:10.1039/D0SC05401A): the probability that AiZynthFinder can find a synthesis
route to a molecule. 0 = hard / no route found, 1 = readily synthesizable.

Also returns SCScore (Coley 2018 synthetic complexity, 1 easy - 5 hard).
POST /score  {"smiles": ["...", ...]}  ->  {"results": [{"smiles","RAScore","SCScore"}, ...]}
CORS-enabled, no API key. Deploy as an HF Space (Docker SDK) or local Docker.
"""
from fastapi import FastAPI, HTTPException, Response
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel
from typing import List
import json, logging, os

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

app = FastAPI(title="RAScore API", description="Retrosynthetic accessibility score for MolParetoLab", version="1.0.0")
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=False,
    allow_methods=["GET", "POST", "OPTIONS"],
    allow_headers=["*"],
)

_scorer = None

def get_scorer():
    global _scorer
    if _scorer is None:
        logger.info("Loading RAScore XGB model...")
        from RAscore import RAscore_XGB
        _scorer = RAscore_XGB.RAScorerXGB()
        logger.info("RAScore model loaded.")
    return _scorer

def ra_score(smiles: str):
    try:
        from rdkit import Chem
        if Chem.MolFromSmiles(smiles) is None:
            return None
        return round(float(get_scorer().predict(smiles)), 4)
    except Exception:
        import traceback
        logger.error(f"RAScore failed for {smiles}:\n{traceback.format_exc()}")
        return None


# ── SA score: Ertl & Schuffenhauer 2009 (synthetic accessibility, 1 easy – 10 hard),
#    via RDKit's Contrib sascorer — so this one endpoint returns all three make-ability scores.
_sascorer = None

def sa_score(smiles: str):
    global _sascorer
    try:
        if _sascorer is None:
            import sys, os as _os
            from rdkit.Chem import RDConfig
            sys.path.append(_os.path.join(RDConfig.RDContribDir, "SA_Score"))
            import sascorer as _sa
            _sascorer = _sa
        from rdkit import Chem
        m = Chem.MolFromSmiles(smiles)
        return round(_sascorer.calculateScore(m), 2) if m is not None else None
    except Exception:
        import traceback
        logger.error(f"SA score failed for {smiles}:\n{traceback.format_exc()}")
        return None


# ── SCScore: Coley et al., J. Chem. Inf. Model. 2018 (synthetic complexity, 1 easy – 5 hard).
#    Standalone numpy model (no TensorFlow) + 1024-bit folded-FP weights, fetched in the image.
_scscorer = None

def get_scscorer():
    global _scscorer
    if _scscorer is None:
        import sys
        if "/app" not in sys.path:
            sys.path.append("/app")
        logger.info("Loading SCScore model...")
        from scscore_standalone import SCScorer
        s = SCScorer()
        s.restore("/app/scscore_1024bool.json.gz", FP_rad=2, FP_len=1024)
        _scscorer = s
        logger.info("SCScore model loaded.")
    return _scscorer

def sc_score(smiles: str):
    try:
        _, score = get_scscorer().get_score_from_smi(smiles)
        return round(float(score), 3)
    except Exception:
        import traceback
        logger.error(f"SCScore failed for {smiles}:\n{traceback.format_exc()}")
        return None


class ScoreRequest(BaseModel):
    smiles: List[str]


@app.get("/health")
def health():
    return {"status": "ok", "model_loaded": _scorer is not None}


@app.post("/score")
def score(req: ScoreRequest):
    if not req.smiles:
        raise HTTPException(400, "smiles list is empty")
    if len(req.smiles) > 1000:
        raise HTTPException(400, "Maximum 1000 SMILES per request")
    results = [{"smiles": s, "SA_Score": sa_score(s), "RAScore": ra_score(s), "SCScore": sc_score(s)} for s in req.smiles]
    return Response(content=json.dumps({"results": results}), media_type="application/json")


if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=int(os.environ.get("PORT", 7860)))