# © 四川森普力科技有限公司  联系人：姚工  电话：159-2808-9752/177-8067-0842
# 邮件：mylanview@126.com   网址：www.sploptics.com
"""E2 — 物质类型鉴定插件（光谱相似度比对）"""

import json
import sqlite3
import uuid
from dataclasses import dataclass, field
from datetime import datetime
from pathlib import Path
from typing import Optional

import numpy as np

_HERE = Path(__file__).parent

try:
    from plugins.base_plugin import PluginInterface, PluginInfo, PluginInput, PluginOutput  # type: ignore
except ImportError:
    import sys
    sys.path.insert(0, str(_HERE.parent.parent))
    from plugins.base_plugin import PluginInterface, PluginInfo, PluginInput, PluginOutput


@dataclass
class SpectrumRecord:
    record_id: str
    name: str
    category: str
    data_type: str
    wavelengths: np.ndarray
    intensities: np.ndarray
    notes: str = ""
    source: str = ""


@dataclass
class MaterialDatabase:
    db_path: str
    db_name: str
    record_count: int
    categories: list
    wl_range: tuple


@dataclass
class MaterialMatch:
    record: SpectrumRecord
    similarity: float
    cosine_sim: float
    pearson_r: float
    rank: int
    overlap_nm: float


# ── 内存缓存 ─────────────────────────────────────────────────────────────

_db_cache: dict = {}   # db_path → list[SpectrumRecord]


def _init_db(db_path: str) -> None:
    """初始化空 SQLite 数据库（表结构）。"""
    conn = sqlite3.connect(db_path)
    conn.execute("""
        CREATE TABLE IF NOT EXISTS spectra (
            id TEXT PRIMARY KEY,
            name TEXT NOT NULL,
            category TEXT NOT NULL DEFAULT '未分类',
            data_type TEXT NOT NULL,
            wavelengths TEXT NOT NULL,
            intensities TEXT NOT NULL,
            notes TEXT DEFAULT '',
            source TEXT DEFAULT '',
            created_at TEXT NOT NULL
        )
    """)
    conn.commit()
    conn.close()


def _load_records(db_path: str) -> list:
    conn = sqlite3.connect(db_path)
    rows = conn.execute(
        "SELECT id, name, category, data_type, wavelengths, intensities, notes, source FROM spectra"
    ).fetchall()
    conn.close()
    records = []
    for row in rows:
        records.append(SpectrumRecord(
            record_id=row[0], name=row[1], category=row[2], data_type=row[3],
            wavelengths=np.array(json.loads(row[4]), dtype=float),
            intensities=np.array(json.loads(row[5]), dtype=float),
            notes=row[6], source=row[7],
        ))
    return records


def _load_records_from_csv(csv_path: str) -> list:
    """从 CSV 文件加载谱线记录（首列波长，其余列物质）。"""
    import csv
    records = []
    with open(csv_path, newline="", encoding="utf-8-sig") as f:
        reader = csv.reader(f)
        header = next(reader)
        rows = list(reader)
    if not rows:
        return records
    wl_col = np.array([float(r[0]) for r in rows], dtype=float)
    names = header[1:]
    for col_idx, name in enumerate(names):
        intensities = np.array([float(r[col_idx + 1]) for r in rows], dtype=float)
        records.append(SpectrumRecord(
            record_id=str(uuid.uuid4()),
            name=name, category="未分类", data_type="reflectance",
            wavelengths=wl_col.copy(), intensities=intensities,
        ))
    return records


def load_database(db_path: str) -> MaterialDatabase:
    """加载数据库（.db SQLite 或 .csv）。同路径内存缓存，不重复IO。"""
    path = Path(db_path)

    if db_path not in _db_cache:
        if path.suffix.lower() == ".csv":
            if not path.exists():
                raise FileNotFoundError(f"CSV 数据库文件不存在：{db_path}")
            _db_cache[db_path] = _load_records_from_csv(db_path)
        else:
            if not path.exists():
                _init_db(db_path)
            _db_cache[db_path] = _load_records(db_path)

    records = _db_cache[db_path]
    categories = sorted(set(r.category for r in records))
    wl_all = [r.wavelengths for r in records]
    if wl_all:
        wl_range = (float(min(w[0] for w in wl_all)), float(max(w[-1] for w in wl_all)))
    else:
        wl_range = (0.0, 0.0)

    return MaterialDatabase(
        db_path=db_path,
        db_name=path.stem,
        record_count=len(records),
        categories=categories,
        wl_range=wl_range,
    )


def _cosine_similarity(a: np.ndarray, b: np.ndarray) -> float:
    denom = np.linalg.norm(a) * np.linalg.norm(b)
    return float(np.dot(a, b) / denom) if denom > 0 else 0.0


def _pearson_r(a: np.ndarray, b: np.ndarray) -> float:
    if len(a) < 2:
        return 0.0
    corr = np.corrcoef(a, b)
    return float(corr[0, 1]) if not np.isnan(corr[0, 1]) else 0.0


def _minmax(arr: np.ndarray) -> np.ndarray:
    lo, hi = arr.min(), arr.max()
    return (arr - lo) / (hi - lo + 1e-9)


def identify(
    wl: np.ndarray,
    intensities: np.ndarray,
    db: MaterialDatabase,
    top_n: int = 5,
    preprocess: bool = True,
    category_filter: Optional[str] = None,
) -> list:
    """相似度鉴定，返回 top_n 结果（按 similarity 降序）。"""
    records = _db_cache.get(db.db_path, [])
    if category_filter:
        records = [r for r in records if r.category == category_filter]

    results = []
    for rec in records:
        # 计算重叠范围
        lo = max(float(wl[0]), float(rec.wavelengths[0]))
        hi = min(float(wl[-1]), float(rec.wavelengths[-1]))
        overlap = hi - lo
        if overlap < 50.0:
            continue

        # 插值对齐到 ref 波长网格
        mask_q = (wl >= lo) & (wl <= hi)
        wl_q = wl[mask_q]
        int_q = intensities[mask_q]
        int_ref = np.interp(wl_q, rec.wavelengths, rec.intensities)

        if preprocess:
            int_q = _minmax(int_q)
            int_ref = _minmax(int_ref)

        cos = _cosine_similarity(int_q, int_ref)
        pear = _pearson_r(int_q, int_ref)
        sim = 0.5 * cos + 0.5 * pear

        results.append(MaterialMatch(
            record=rec, similarity=sim,
            cosine_sim=cos, pearson_r=pear,
            rank=0, overlap_nm=overlap,
        ))

    results.sort(key=lambda m: m.similarity, reverse=True)
    for i, m in enumerate(results[:top_n]):
        m.rank = i + 1
    return results[:top_n]


def import_from_csv(
    csv_path: str,
    db_path: str,
    category: str = "未分类",
    data_type: str = "reflectance",
    force_overwrite: bool = False,
) -> int:
    """从 CSV 批量导入谱线。首列波长，其余列物质。返回新增记录数。"""
    import csv
    path = Path(db_path)
    if not path.exists():
        _init_db(db_path)

    rows = []
    with open(csv_path, newline="", encoding="utf-8-sig") as f:
        reader = csv.reader(f)
        header = next(reader)
        for row in reader:
            rows.append(row)

    if not rows:
        return 0

    wl_col = [float(r[0]) for r in rows]
    if wl_col != sorted(wl_col):
        raise ValueError("波长列须单调递增")

    names = header[1:]
    conn = sqlite3.connect(db_path)
    count = 0
    for col_idx, name in enumerate(names):
        intensities = [float(r[col_idx + 1]) for r in rows]
        existing = conn.execute(
            "SELECT id FROM spectra WHERE name=? AND category=?", (name, category)
        ).fetchone()
        if existing:
            if force_overwrite:
                conn.execute(
                    "UPDATE spectra SET data_type=?, wavelengths=?, intensities=?, created_at=? WHERE id=?",
                    (data_type, json.dumps(wl_col), json.dumps(intensities),
                     datetime.now().isoformat(), existing[0]),
                )
                count += 1
            # force_overwrite=False 时跳过同名记录
        else:
            rid = str(uuid.uuid4())
            conn.execute(
                "INSERT INTO spectra VALUES (?,?,?,?,?,?,?,?,?)",
                (rid, name, category, data_type,
                 json.dumps(wl_col), json.dumps(intensities),
                 "", "", datetime.now().isoformat()),
            )
            count += 1
    conn.commit()
    conn.close()

    # 清除缓存
    _db_cache.pop(db_path, None)
    return count


def export_database(db: MaterialDatabase, export_dir: str, format: str = "csv") -> str:
    """导出数据库为 CSV。"""
    import csv
    records = _db_cache.get(db.db_path, [])
    out_path = Path(export_dir) / f"{db.db_name}_export.csv"
    with open(str(out_path), "w", newline="", encoding="utf-8-sig") as f:
        writer = csv.writer(f)
        if records:
            wl_ref = records[0].wavelengths
            writer.writerow(["wavelength"] + [r.name for r in records])
            for i, wl_val in enumerate(wl_ref):
                row = [wl_val] + [
                    float(np.interp(wl_val, r.wavelengths, r.intensities))
                    for r in records
                ]
                writer.writerow(row)
    return str(out_path)


class MaterialIdPlugin(PluginInterface):

    def __init__(self):
        self._current_db: Optional[MaterialDatabase] = None
        default_db = str(_HERE / "databases" / "default.db")
        try:
            self._current_db = load_database(default_db)
        except Exception:
            pass

    def get_info(self) -> PluginInfo:
        return PluginInfo(
            plugin_id="material_id_v1",
            name="物质类型鉴定",
            version="1.0.0",
            author="四川森普力科技有限公司",
            description="加载用户光谱数据库，与输入谱进行相似度比对，输出匹配物质排名",
            input_type="spectrum",
            output_type="json",
            requires_stop=True,
        )

    def validate_input(self, plugin_input: PluginInput) -> tuple:
        if self._current_db is None:
            return False, "数据库未加载，请先加载光谱数据库"
        if len(plugin_input.wl) < 10:
            return False, "数据点数不足（需 ≥ 10）"
        wl = plugin_input.wl
        lo = max(float(wl[0]), self._current_db.wl_range[0])
        hi = min(float(wl[-1]), self._current_db.wl_range[1])
        if hi - lo < 50.0:
            return False, "输入谱与数据库重叠不足50nm"
        return True, ""

    def run(self, plugin_input: PluginInput) -> PluginOutput:
        matches = identify(plugin_input.wl, plugin_input.intensities, self._current_db)

        result_list = []
        for m in matches:
            result_list.append({
                "rank": m.rank,
                "name": m.record.name,
                "category": m.record.category,
                "similarity": round(m.similarity, 4),
                "cosine_sim": round(m.cosine_sim, 4),
                "pearson_r": round(m.pearson_r, 4),
                "overlap_nm": round(m.overlap_nm, 1),
            })

        if matches:
            top = matches[0]
            summary = (f"最佳匹配：{top.record.name}（{top.record.category}）  "
                       f"相似度={top.similarity:.2%}  重叠={top.overlap_nm:.0f}nm")
        else:
            summary = "未找到匹配（数据库为空或重叠不足）"

        return PluginOutput(
            plugin_id="material_id_v1",
            result={"matches": result_list, "db_name": self._current_db.db_name
                    if self._current_db else ""},
            summary=summary,
        )
