"""Reconstruct DNA from raw mean-pooled embeddings saved as a NumPy array."""

import argparse
import json
from pathlib import Path

import numpy as np
import torch
from transformers import AutoTokenizer

from model import QueryDecoderReconstructor


class Decoder:
    def __init__(self, model_dir, device="cpu"):
        folder = Path(model_dir)
        checkpoint = torch.load(folder / "model.pt", map_location="cpu", weights_only=True)
        config = checkpoint["config"]["model"]
        self.device = torch.device(device)
        self.input_dim = checkpoint["input_dim"]
        self.stats = json.loads((folder / "normalization.json").read_text())
        assert self.stats["method"] == "standard"
        assert np.isfinite([self.stats["mean"], self.stats["std"]]).all()
        assert self.stats["std"] > 0
        self.model = QueryDecoderReconstructor(
            input_dim=self.input_dim, mode=checkpoint["mode"],
            seq_length=checkpoint["effective_seq_length"], output_dim=checkpoint["output_dim"],
            **{key: config[key] for key in ["d_model", "nhead", "num_layers", "dim_feedforward",
                                           "dropout", "n_context_tokens", "aux_length_loss_weight"]},
        )
        self.model.load_state_dict(checkpoint["state_dict"], strict=True)
        self.model.to(self.device).eval()
        self.character = checkpoint["tokenizer_type"] == "char"
        if self.character:
            self.eos_id = 4
            assert checkpoint["output_dim"] == 5
        else:
            self.tokenizer = AutoTokenizer.from_pretrained(folder / "tokenizer", local_files_only=True)
            self.eos_id = json.loads((folder / "tokenizer_info.json").read_text())["eos_id"]
            assert len(self.tokenizer) == checkpoint["output_dim"]

    def decode_tokens(self, ids):
        ids = list(ids)
        if self.eos_id in ids:
            ids = ids[:ids.index(self.eos_id)]
        if self.character:
            return "".join("ACGT"[i] for i in ids)
        return self.tokenizer.decode(ids, skip_special_tokens=True).replace(" ", "")

    @torch.inference_mode()
    def predict(self, embeddings, batch_size=32):
        raw = np.asarray(embeddings, dtype=np.float32)
        assert raw.ndim == 2 and raw.shape[1] == self.input_dim, f"Expected [N, {self.input_dim}] embeddings"
        assert np.isfinite(raw).all(), "Embeddings must be finite"
        assert batch_size > 0
        sequences = []
        for start in range(0, len(raw), batch_size):
            normalized = (raw[start:start + batch_size] - self.stats["mean"]) / self.stats["std"]
            assert np.isfinite(normalized).all()
            logits, _ = self.model(torch.from_numpy(normalized).to(self.device))
            assert torch.isfinite(logits).all(), "Non-finite model output"
            sequences.extend(self.decode_tokens(ids) for ids in logits.argmax(-1).cpu().tolist())
        return sequences


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--model", required=True, choices=["dnabert2", "evo2", "ntv2"])
    parser.add_argument("--embeddings", required=True, type=Path)
    parser.add_argument("--output", required=True, type=Path)
    parser.add_argument("--device", default="cpu")
    parser.add_argument("--batch-size", type=int, default=32)
    args = parser.parse_args()
    decoder = Decoder(Path(__file__).parent / args.model, args.device)
    sequences = decoder.predict(np.load(args.embeddings, allow_pickle=False), args.batch_size)
    args.output.write_text(json.dumps(sequences, indent=2) + "\n", encoding="utf-8")
    print(f"Wrote {len(sequences)} reconstructions to {args.output}")


if __name__ == "__main__":
    main()