from dataclasses import dataclass
from pathlib import Path
from typing import Iterator

import requests
from Bio import SeqIO, Entrez
from Bio.Seq import Seq


Entrez.email = "your@email.example"  # NCBI 필수


@dataclass
class ExonRegion:
    """유전자 exon 하나의 정보."""
    gene_name: str
    exon_number: int
    chromosome: str
    strand: str            # "+" or "-"
    start_hg38: int         # 0-based
    end_hg38: int
    length: int
    sequence: str
    coding_frame: int | None  # None if UTR, 0/1/2 otherwise


def fetch_refseq_mrna(refseq_id: str = "NM_007294") -> str:
    """NCBI에서 RefSeq mRNA 서열 fetch (예: BRCA1 = NM_007294)."""
    with Entrez.efetch(db="nucleotide", id=refseq_id, rettype="fasta", retmode="text") as h:
        record = SeqIO.read(h, "fasta")
    return str(record.seq).upper()


def fetch_gene_exons_ucsc(gene_symbol: str, assembly: str = "hg38") -> list[ExonRegion]:
    """UCSC Genome Browser REST API로 유전자 exon 좌표 fetch.
    
    실전에서는 GENCODE annotation GTF 로컬 파싱이 훨씬 안정적.
    """
    # UCSC Table Browser 대안: MyGene.info API
    resp = requests.get(
        f"https://mygene.info/v3/query?q=symbol:{gene_symbol}&species=human&fields=exons",
        timeout=30,
    )
    resp.raise_for_status()
    hits = resp.json().get("hits", [])
    if not hits:
        return []
    
    exons_info = hits[0].get("exons", [])
    if not exons_info:
        return []
    
    canonical = exons_info[0]  # canonical transcript
    exon_regions = []
    for i, (start, end) in enumerate(canonical.get("position", [])):
        # UCSC에서 실제 서열 fetch (DAS API 또는 로컬 BSgenome)
        seq = _fetch_ucsc_dna(canonical["chr"], start, end, assembly)
        exon_regions.append(ExonRegion(
            gene_name=gene_symbol,
            exon_number=i + 1,
            chromosome=canonical["chr"],
            strand=canonical.get("strand", "+"),
            start_hg38=start, end_hg38=end,
            length=end - start,
            sequence=seq,
            coding_frame=None,  # 상세 판정은 CDS 좌표 필요
        ))
    return exon_regions


def _fetch_ucsc_dna(chrom: str, start: int, end: int, assembly: str = "hg38") -> str:
    """UCSC DAS API로 특정 영역 서열 반환."""
    url = f"https://api.genome.ucsc.edu/getData/sequence?genome={assembly};chrom={chrom};start={start};end={end}"
    resp = requests.get(url, timeout=30)
    resp.raise_for_status()
    return resp.json().get("dna", "").upper()

@dataclass
class GRNACandidate:
    gene_name: str
    exon_number: int
    protospacer: str        # 20bp 표적 서열
    pam: str                # 3bp PAM
    strand: str             # "+" or "-"
    exon_position: int      # exon 서열 상 시작 위치 (0-based)
    genome_position: int    # hg38 게놈 좌표 (0-based)
    context_50bp: str       # 앞 15bp + gRNA + PAM + 뒤 15bp = 53bp (CRISPRon 등 학습 입력)


def scan_pam_ngg(sequence: str, gene_name: str, exon: ExonRegion) -> list[GRNACandidate]:
    """양쪽 가닥에서 NGG PAM + 20bp protospacer 후보 추출."""
    candidates = []
    seq_plus = sequence.upper()
    seq_minus = str(Seq(seq_plus).reverse_complement())
    
    def _scan(strand_seq: str, strand_label: str, seq_len: int) -> None:
        for i in range(20, seq_len - 3):
            pam = strand_seq[i:i+3]
            if pam[1:3] != "GG":
                continue
            protospacer = strand_seq[i-20:i]
            if "N" in protospacer:
                continue
            context_start = max(0, i - 35)
            context_end = min(seq_len, i + 18)
            context = strand_seq[context_start:context_end]
            # 게놈 좌표 계산 (strand별 방향 반전)
            if strand_label == "+":
                genome_pos = exon.start_hg38 + (i - 20)
            else:
                genome_pos = exon.end_hg38 - i
            candidates.append(GRNACandidate(
                gene_name=gene_name,
                exon_number=exon.exon_number,
                protospacer=protospacer,
                pam=pam,
                strand=strand_label,
                exon_position=i - 20,
                genome_position=genome_pos,
                context_50bp=context,
            ))
    
    _scan(seq_plus, "+", len(seq_plus))
    _scan(seq_minus, "-", len(seq_minus))
    return candidates


def scan_gene_all_exons(exons: list[ExonRegion]) -> list[GRNACandidate]:
    all_cands = []
    for exon in exons:
        all_cands.extend(scan_pam_ngg(exon.sequence, exons[0].gene_name, exon))
    return all_cands

from Bio.SeqUtils import GC


def check_self_complementarity(seq: str, min_stem: int = 4) -> bool:
    """자체 상보성 검사: 5' end 4bp가 3' end 4bp의 reverse complement이면 hairpin 우려."""
    if len(seq) < min_stem * 2:
        return False
    stem_5 = seq[:min_stem]
    stem_3 = seq[-min_stem:]
    return stem_5 == str(Seq(stem_3).reverse_complement())


def filter_grna_physicochemical(
    candidates: list[GRNACandidate],
    min_gc: float = 30.0,
    max_gc: float = 70.0,
    max_poly_t: int = 3,
    prefer_g_start: bool = True,
    check_hairpin: bool = True,
) -> list[GRNACandidate]:
    """물리화학 필터."""
    filtered = []
    for c in candidates:
        # GC content
        gc = GC(c.protospacer)
        if not (min_gc <= gc <= max_gc):
            continue
        # Poly-T
        if "T" * (max_poly_t + 1) in c.protospacer:
            continue
        # G start (선호)
        if prefer_g_start and c.protospacer[0] != "G":
            # 완전 제외는 아니지만 후순위 (여기서는 통과)
            pass
        # Hairpin
        if check_hairpin and check_self_complementarity(c.protospacer):
            continue
        filtered.append(c)
    return filtered

import numpy as np
import rpy2.robjects as ro
from rpy2.robjects.packages import importr
from rpy2.robjects.vectors import StrVector


class CrisprScoreEnsemble:
    """R crisprScore의 여러 스코어러를 Python에서 호출."""

    def __init__(self):
        # R 패키지 로딩 (사전에 R 콘솔에서 BiocManager::install 필요)
        self.crisprscore = importr("crisprScore")
        self.base = importr("base")

    def crispron(self, contexts: list[str]) -> list[float]:
        """CRISPRon 온타겟 스코어. 입력: 각 항목 30bp context (4bp upstream + 20bp protospacer + 3bp PAM + 3bp downstream)."""
        r_input = StrVector(contexts)
        try:
            result = self.crisprscore.getCRISPRonScores(r_input)
            return list(result)
        except Exception as e:
            print(f"CRISPRon 실패: {e}")
            return [float("nan")] * len(contexts)

    def deep_hf(self, protospacers_with_pam: list[str], enzyme: str = "WT") -> list[float]:
        """DeepHF 스코어. 입력: 각 23bp (20 protospacer + 3 PAM). enzyme: 'WT', 'ESP', 'HF'."""
        r_input = StrVector(protospacers_with_pam)
        try:
            result = self.crisprscore.getDeepHFScores(r_input, enzyme=enzyme)
            return list(result)
        except Exception as e:
            print(f"DeepHF 실패: {e}")
            return [float("nan")] * len(protospacers_with_pam)

    def crispron_be(self, contexts: list[str], editor: str = "ABE8e") -> list[float]:
        """CRISPRon-BE base editing 효율.
        
        editor 옵션: 'ABE8e' (adenine), 'BE4max' (cytosine), etc.
        """
        r_input = StrVector(contexts)
        try:
            result = self.crisprscore.getCRISPRonBEScores(r_input, editor=editor)
            return list(result)
        except Exception as e:
            print(f"CRISPRon-BE 실패: {e}")
            return [float("nan")] * len(contexts)

    def cfd_off_target(self, protospacer: str, target_dna: str) -> float:
        """Cutting Frequency Determination (Doench 2016 [5])."""
        try:
            result = self.crisprscore.getCFDScores(
                StrVector([protospacer]),
                StrVector([target_dna]),
            )
            return float(list(result)[0])
        except Exception:
            return 0.0

    def mit_hsu_zhang(self, protospacer: str, target_dna: str) -> float:
        """MIT/Hsu-Zhang 오프타겟 스코어 (position-dependent mismatch penalty)."""
        try:
            result = self.crisprscore.getMITScores(
                StrVector([protospacer]),
                StrVector([target_dna]),
            )
            return float(list(result)[0])
        except Exception:
            return 0.0

# R 콘솔에서
if (!require("BiocManager", quietly = TRUE))
    install.packages("BiocManager")
BiocManager::install(c(
    "crisprScore", "crisprDesign",
    "BSgenome.Hsapiens.UCSC.hg38",
    "crisprBase"
))

def query_crispor_offtargets(
    protospacer: str,
    genome: str = "hg38",
    max_mismatches: int = 4,
    timeout: int = 300,
) -> list[dict]:
    """CRISPOR 원격 API로 게놈 전체 오프타겟 스캔.
    
    Returns: [{genome_pos, mismatches, cfd_score, gene}, ...]
    """
    # 실전 endpoint: http://crispor.tefor.net/crispor.py
    # 여기서는 개념적 wrapper. 실제 API는 HTML 반환이라 파싱 필요.
    try:
        resp = requests.get(
            "http://crispor.tefor.net/crispor.py",
            params={
                "seq": protospacer,
                "org": genome,
                "pam": "NGG",
                "showAllOTs": "1",
            },
            timeout=timeout,
        )
        resp.raise_for_status()
        # 실제 CRISPOR 응답은 HTML/TSV. 여기서는 예시 파싱만 표현.
        # 실전에서는 BeautifulSoup 또는 crispor 로컬 CLI로 TSV 획득 권장.
        return _parse_crispor_html(resp.text)
    except Exception as e:
        print(f"CRISPOR 오프타겟 스캔 실패 ({protospacer}): {e}")
        return []


def _parse_crispor_html(html: str) -> list[dict]:
    """CRISPOR HTML 응답 파싱 (개념 스텁, 실전은 BeautifulSoup)."""
    return []  # 실전 구현 시 완성


def compute_off_target_summary(off_targets: list[dict]) -> dict:
    """오프타겟 스캔 결과 요약."""
    if not off_targets:
        return {
            "n_offtargets_mm4": 0,
            "max_cfd_score": 0.0,
            "high_risk_genes": [],
            "genome_safety_score": 1.0,
        }
    max_cfd = max(ot["cfd_score"] for ot in off_targets)
    high_risk = [ot["gene"] for ot in off_targets if ot["cfd_score"] > 0.5 and ot["gene"]]
    # 종합 안전성: 최대 CFD 뒤집기 (0=위험, 1=안전)
    safety = 1.0 - max_cfd
    return {
        "n_offtargets_mm4": len(off_targets),
        "max_cfd_score": max_cfd,
        "high_risk_genes": high_risk[:5],
        "genome_safety_score": safety,
    }

import pandas as pd


DEFAULT_WEIGHTS = {
    "crispron": 0.30,
    "deep_hf": 0.25,
    "crispron_be": 0.10,  # BE 시나리오 아닐 때는 0으로
    "off_target_safety": 0.35,
}


def build_ensemble_ranking(
    candidates: list[GRNACandidate],
    on_scores: dict[str, list[float]],
    off_target_summaries: list[dict],
    weights: dict[str, float] = DEFAULT_WEIGHTS,
) -> pd.DataFrame:
    """스코어러 통합 랭킹 DataFrame."""
    df = pd.DataFrame([{
        "gene": c.gene_name,
        "exon": c.exon_number,
        "protospacer": c.protospacer,
        "pam": c.pam,
        "strand": c.strand,
        "genome_position": c.genome_position,
        "context_50bp": c.context_50bp,
    } for c in candidates])
    
    for name, scores in on_scores.items():
        df[f"on_{name}"] = scores
    
    df["off_n_mm4"] = [s["n_offtargets_mm4"] for s in off_target_summaries]
    df["off_max_cfd"] = [s["max_cfd_score"] for s in off_target_summaries]
    df["off_safety"] = [s["genome_safety_score"] for s in off_target_summaries]
    df["off_high_risk_genes"] = [
        ";".join(s["high_risk_genes"]) for s in off_target_summaries
    ]
    
    # 정규화 후 가중 평균
    def _normalize(col):
        c = df[col].fillna(df[col].median())
        return (c - c.min()) / (c.max() - c.min() + 1e-8)
    
    final = np.zeros(len(df))
    if "on_crispron" in df.columns:
        final += weights.get("crispron", 0.0) * _normalize("on_crispron")
    if "on_deep_hf" in df.columns:
        final += weights.get("deep_hf", 0.0) * _normalize("on_deep_hf")
    if "on_crispron_be" in df.columns:
        final += weights.get("crispron_be", 0.0) * _normalize("on_crispron_be")
    final += weights["off_target_safety"] * df["off_safety"].fillna(0.5)
    
    df["final_score"] = final
    return df.sort_values("final_score", ascending=False).reset_index(drop=True)

def generate_lab_report(
    ranked: pd.DataFrame,
    output_md: Path,
    output_csv: Path,
    top_k: int = 5,
) -> None:
    """실험실 발주 가능한 리포트 생성."""
    top = ranked.head(top_k)
    
    lines = [
        f"# gRNA 후보 상위 {top_k}개 발주 리포트",
        "",
        f"대상 유전자: **{top.iloc[0]['gene']}**",
        f"생성 시각: (파이프라인 실행 시각으로 채우기)",
        "",
        "## 요약 표",
        "",
        "| Rank | Exon | Protospacer (20bp) | PAM | Strand | On (CRISPRon) | On (DeepHF) | Off Safety | Final |",
        "|:----:|:----:|:------------------:|:---:|:------:|:-------------:|:-----------:|:----------:|:-----:|",
    ]
    for i, row in top.iterrows():
        lines.append(
            f"| {i+1} | {row['exon']} | `{row['protospacer']}` | {row['pam']} | {row['strand']} | "
            f"{row.get('on_crispron', float('nan')):.3f} | {row.get('on_deep_hf', float('nan')):.3f} | "
            f"{row['off_safety']:.3f} | **{row['final_score']:.3f}** |"
        )
    
    lines += [
        "",
        "## 발주 서열 (5' → 3')",
        "",
        "Spacer 서열에 U6 promoter 요구 G가 없으면 앞에 G 추가 or ExoScribe 등 alternative promoter 활용.",
        "",
    ]
    for i, row in top.iterrows():
        spacer = row["protospacer"]
        if spacer[0] != "G":
            forge = f"G{spacer}"
        else:
            forge = spacer
        lines.append(f"### Rank {i+1}")
        lines.append(f"- Protospacer (spacer only, 20bp): `{spacer}`")
        lines.append(f"- Order-ready (with G prefix if needed): `{forge}` ({len(forge)}bp)")
        lines.append(f"- PAM: `{row['pam']}`")
        lines.append(f"- Genome position (hg38): {row['genome_position']}")
        if row["off_high_risk_genes"]:
            lines.append(f"- ⚠️ High-risk off-target genes: {row['off_high_risk_genes']}")
        lines.append("")
    
    lines += [
        "## 사전 실험 체크리스트",
        "",
        "- [ ] 발주 전 alignment 재확인 (BLASTn to hg38 · GRCh38.p14)",
        "- [ ] SpCas9-HF1 vector 준비 (예: pX330 or plentiCRISPR v2)",
        "- [ ] 대상 세포주 (K562)의 mycoplasma 검사 완료",
        "- [ ] T7E1 assay 또는 amplicon deep sequencing primer 설계",
        "- [ ] 대조군 gRNA (non-targeting, e.g. sgLacZ) 병행 발주",
    ]
    
    output_md.write_text("\n".join(lines), encoding="utf-8")
    top.to_csv(output_csv, index=False)

def full_pipeline(
    gene_symbol: str,
    refseq_id: str,
    output_prefix: str,
    editor: str | None = None,  # None: nuclease KO, "ABE8e"/"BE4max": base editing
    top_k: int = 5,
) -> pd.DataFrame:
    print(f"[1/7] {gene_symbol} exon 정보 fetch (RefSeq {refseq_id})")
    # 실전에서는 exon 좌표를 GENCODE GTF에서 로컬 파싱 권장
    exons = fetch_gene_exons_ucsc(gene_symbol)
    if not exons:
        # fallback: mRNA 전체를 단일 "exon"으로 처리 (간단 시연)
        seq = fetch_refseq_mrna(refseq_id)
        exons = [ExonRegion(
            gene_name=gene_symbol, exon_number=1, chromosome="",
            strand="+", start_hg38=0, end_hg38=len(seq),
            length=len(seq), sequence=seq, coding_frame=0,
        )]
    print(f"      exon {len(exons)}개")
    
    print("[2/7] PAM (NGG) 스캔")
    all_cands = scan_gene_all_exons(exons)
    print(f"      원시 후보 {len(all_cands)}")
    
    print("[3/7] 물리·화학 필터")
    filtered = filter_grna_physicochemical(all_cands)
    print(f"      필터 통과 {len(filtered)}")
    
    print("[4/7] 온타겟 스코어러 앙상블")
    scorer = CrisprScoreEnsemble()
    contexts = [c.context_50bp for c in filtered]
    proto_pam = [c.protospacer + c.pam for c in filtered]
    on_scores = {
        "crispron": scorer.crispron(contexts),
        "deep_hf": scorer.deep_hf(proto_pam, enzyme="HF"),  # SpCas9-HF1
    }
    if editor:
        on_scores["crispron_be"] = scorer.crispron_be(contexts, editor=editor)
    
    print("[5/7] 오프타겟 스캔 (CRISPOR 원격)")
    off_summaries = []
    for i, c in enumerate(filtered):
        if i % 50 == 0:
            print(f"      진행 {i}/{len(filtered)}")
        offs = query_crispor_offtargets(c.protospacer)
        off_summaries.append(compute_off_target_summary(offs))
    
    print("[6/7] 앙상블 랭킹")
    ranked = build_ensemble_ranking(filtered, on_scores, off_summaries)
    ranked.to_csv(f"{output_prefix}_all_candidates.csv", index=False)
    
    print(f"[7/7] 상위 {top_k} 발주 리포트")
    generate_lab_report(
        ranked,
        Path(f"{output_prefix}_top{top_k}.md"),
        Path(f"{output_prefix}_top{top_k}.csv"),
        top_k=top_k,
    )
    print(f"      완료: {output_prefix}_top{top_k}.md")
    return ranked


# 실행 예시 (BRCA1 KO 시나리오)
# ranked = full_pipeline(
#     gene_symbol="BRCA1",
#     refseq_id="NM_007294",
#     output_prefix="brca1_k562",
#     editor=None,  # SpCas9-HF1 nuclease KO
#     top_k=5,
# )

[1/7] BRCA1 exon 정보 fetch (RefSeq NM_007294)
      exon 24개
[2/7] PAM (NGG) 스캔
      원시 후보 423
[3/7] 물리·화학 필터
      필터 통과 197
[4/7] 온타겟 스코어러 앙상블
[5/7] 오프타겟 스캔 (CRISPOR 원격)
      진행 0/197
      진행 50/197
      ...
[6/7] 앙상블 랭킹
[7/7] 상위 5 발주 리포트
      완료: brca1_k562_top5.md