import argparse
import os
from pathlib import Path

import numpy as np
import tensorflow as tf

os.environ.setdefault("TF_FORCE_GPU_ALLOW_GROWTH", "true")

AA_ORDER = ["A", "R", "N", "D", "C", "Q", "E", "G", "H", "I",
            "L", "K", "M", "F", "P", "S", "T", "W", "Y", "V", "X", "-"]

def decode_one(img):
    arr = np.asarray(img)
    arr = arr.reshape(32, 22)

    # remove X column, same as README
    arr_no_x = np.delete(arr, 20, axis=1)
    aa_no_x = AA_ORDER[:20] + ["-"]

    aa_idx = np.argmax(arr_no_x, axis=1)
    aa = np.array([aa_no_x[i] for i in aa_idx], dtype=object)

    gaps = np.where(aa == "-")[0] + 1  # use 1-based position to mimic R logic
    if len(gaps) > 0:
        gap_first_candidates = gaps[gaps > 1]
        gap_last_candidates = gaps[gaps < 32]
        if len(gap_first_candidates) > 0 and len(gap_last_candidates) > 0:
            gap_first = int(gap_first_candidates[0])
            gap_last = int(gap_last_candidates[-1])
            if gap_first <= gap_last:
                aa[(gap_first - 1):gap_last] = "-"

    seq = "".join(aa.tolist()).replace("-", "")
    return seq

def generate_one(model_id, n_seq=100, batch_size=20, latent_dim=100, seed=2026):
    #model_dir = Path(f"weight/GAN/GAN_model_{model_id}_dcu.keras")
    model_dir = Path(f"weight/GAN/GAN_model_{model_id}_dcu")
    if not model_dir.exists():
        raise FileNotFoundError(f"Missing trained model: {model_dir}")

    np.random.seed(seed + model_id)
    tf.random.set_seed(seed + model_id)

    model = tf.keras.models.load_model(str(model_dir), compile=False)

    seqs = []
    with tf.device("/GPU:0"):
        while len(seqs) < n_seq:
            noise = np.random.normal(size=(batch_size, latent_dim)).astype("float32")
            fake = model(noise, training=False).numpy()

            for i in range(fake.shape[0]):
                seqs.append(decode_one(fake[i]))
                if len(seqs) >= n_seq:
                    break

    return seqs

def load_group_names():
    names = []
    for i in range(1, 16):
        f = Path(f"model/GAN/seq_encoded_{i:02d}.npz")
        if f.exists():
            d = np.load(f, allow_pickle=True)
            names.append(str(d["name"]))
        else:
            names.append(f"GAN_model_{i}")
    return names

def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--model-id", type=int, default=0, help="0 means all models; otherwise 1-15")
    parser.add_argument("--n-seq", type=int, default=100)
    parser.add_argument("--batch-size", type=int, default=20)
    parser.add_argument("--latent-dim", type=int, default=100)
    parser.add_argument("--seed", type=int, default=2026)
    parser.add_argument("--out-tsv", default="model/GAN/gen_seq_trained_dcu.tsv")
    args = parser.parse_args()

    group_names = load_group_names()

    if args.model_id == 0:
        model_ids = range(1, 16)
    else:
        if args.model_id < 1 or args.model_id > 15:
            raise ValueError("--model-id must be 0 or 1-15")
        model_ids = [args.model_id]

    rows = []
    all_seqs = {}

    for model_id in model_ids:
        group = group_names[model_id - 1]
        print(f"\nGenerating from trained model {model_id}: {group}")
        seqs = generate_one(
            model_id=model_id,
            n_seq=args.n_seq,
            batch_size=args.batch_size,
            latent_dim=args.latent_dim,
            seed=args.seed,
        )
        all_seqs[model_id] = seqs

        unique = []
        seen = set()
        for s in seqs:
            if s not in seen:
                unique.append(s)
                seen.add(s)

        print("unique first 10:")
        print(unique[:10])
        print("n_seq:", len(seqs), "n_unique:", len(unique))

        for rank, seq in enumerate(seqs, start=1):
            rows.append((model_id, group, rank, seq, len(seq)))

    out_path = Path(args.out_tsv)
    out_path.parent.mkdir(parents=True, exist_ok=True)

    with out_path.open("w") as f:
        f.write("model_id\tgroup\trank\taa\tlength\n")
        for row in rows:
            f.write("\t".join(map(str, row)) + "\n")

    np.savez_compressed(
        str(out_path).replace(".tsv", ".npz"),
        **{f"model_{k:02d}": np.array(v, dtype=object) for k, v in all_seqs.items()}
    )

    print("\nsaved:", out_path)
    print("saved:", str(out_path).replace(".tsv", ".npz"))
    print("GAN trained-model generation OK")

if __name__ == "__main__":
    main()