#!/usr/bin/env python
from __future__ import annotations

import argparse
import json
from pathlib import Path

from assaylens.data.access import iter_limited_rows
from assaylens.data.hf_upload import auto_upload_folder
from assaylens.data.schemas import EVE_HF_ID, GOA_HF_ID
from assaylens.features.assay_frame import assay_frame_features
from assaylens.features.chemistry import featurize_smiles
from assaylens.features.goa import build_goa_count_features
from assaylens.features.target import target_features


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--limit", type=int, default=None)
    parser.add_argument("--goa-limit", type=int, default=None)
    parser.add_argument("--fingerprint-bits", type=int, default=256)
    parser.add_argument("--out-dir", default="/tmp/assaylens_outputs/features")
    args = parser.parse_args()
    rows = list(iter_limited_rows(EVE_HF_ID, "train", args.limit))
    goa_rows = list(iter_limited_rows(GOA_HF_ID, "train", args.goa_limit))
    goa_features = build_goa_count_features(goa_rows)
    compound_features_by_id = {}
    target_features_by_id = {}
    assay_feature_rows = []
    for row in rows:
        compound_id = row.get("compound_id")
        if compound_id not in compound_features_by_id:
            chemistry = featurize_smiles(row.get("compound__smiles"), n_bits=args.fingerprint_bits)
            compound_features_by_id[compound_id] = {
                "compound_id": compound_id,
                "compound__smiles": row.get("compound__smiles"),
                **{key: value for key, value in chemistry.items() if key != "fingerprint"},
                "fingerprint": chemistry["fingerprint"],
            }
        target_id = row.get("target_id")
        if target_id not in target_features_by_id:
            target_features_by_id[target_id] = {"target_id": target_id, **target_features(row)}
        assay_feature_rows.append(
            {
                "assay_id": row.get("assay_id"),
                "target_id": row.get("target_id"),
                "compound_id": row.get("compound_id"),
                **assay_frame_features(row),
            }
        )
    out = Path(args.out_dir)
    out.mkdir(parents=True, exist_ok=True)
    import pandas as pd

    pd.DataFrame(list(compound_features_by_id.values())).to_parquet(
        out / "compound_features.parquet", index=False
    )
    pd.DataFrame(list(target_features_by_id.values())).to_parquet(
        out / "target_features.parquet", index=False
    )
    pd.DataFrame(assay_feature_rows).to_parquet(out / "assay_frame_features.parquet", index=False)
    pd.DataFrame(
        [{"target_key": key, **value} for key, value in goa_features.items()]
    ).to_parquet(out / "goa_features.parquet", index=False)
    (out / "features_sample.json").write_text(
        json.dumps(
            {
                "n_rows": len(rows),
                "n_compounds": len(compound_features_by_id),
                "n_targets": len(target_features_by_id),
                "n_goa_keys": len(goa_features),
            },
            indent=2,
        ),
        encoding="utf-8",
    )
    uploaded = auto_upload_folder(out, "features")
    if uploaded:
        print(f"uploaded feature artifacts to {uploaded}/features")
    print(f"wrote feature artifacts to {out}")


if __name__ == "__main__":
    main()