#!/usr/bin/env python3
"""
prepare_cxr_from_single_json.py

Utility to split a single JSON/JSONL dataset into train/valid and prepare
the minimal files expected by the CXR-only training path in RadIR:

- train.jsonl / valid.jsonl : each line JSON with keys `name` and `img_path`
- csv/all.csv : contains a `File Path` column listing sample names (train set)
- npy/all.npy : similarity matrix (float32, values in [0,1]) for the CSV rows

By default similarity is computed with TF-IDF + cosine similarity. If you
have `sentence-transformers` installed you can pass `--sim-method st` to use
a semantic embedding model (better but requires the package and a model).

Examples:
  python scripts/prepare_cxr_from_single_json.py \
    --input dataset/all.jsonl \
    --out /tmp/radir_debug --split 0.9 --seed 42 --sim-method tfidf

"""
import os
import sys
import json
import argparse
from pathlib import Path
import numpy as np


def read_input(path):
    path = Path(path)
    text = path.read_text()
    # Try JSONL (one json per line)
    try:
        lines = [l.strip() for l in text.splitlines() if l.strip()]
        data = [json.loads(l) for l in lines]
        return data
    except Exception:
        pass
    # Try a single JSON array
    try:
        data = json.loads(text)
        if isinstance(data, list):
            return data
    except Exception:
        pass
    raise ValueError(f"Unsupported input format for {path}. Expect JSONL or JSON array of objects.")


def ensure_name_field(entry):
    if 'name' in entry:
        return entry['name']
    if 'img_path' in entry:
        return os.path.basename(entry['img_path'])
    raise ValueError('Entry lacks both `name` and `img_path`')


def save_jsonl(entries, path):
    p = Path(path)
    p.parent.mkdir(parents=True, exist_ok=True)
    with p.open('w') as f:
        for e in entries:
            f.write(json.dumps(e, ensure_ascii=False) + '\n')


def save_csv_filepaths(names, path):
    import csv
    p = Path(path)
    p.parent.mkdir(parents=True, exist_ok=True)
    with p.open('w', newline='') as f:
        writer = csv.writer(f)
        writer.writerow(['File Path'])
        for n in names:
            writer.writerow([n])


def compute_similarity_tfidf(texts):
    try:
        from sklearn.feature_extraction.text import TfidfVectorizer
        from sklearn.metrics.pairwise import cosine_similarity
    except Exception as e:
        raise RuntimeError('TF-IDF dependencies missing: install scikit-learn') from e
    vec = TfidfVectorizer(max_features=8192)
    X = vec.fit_transform(texts)
    sim = cosine_similarity(X)
    # clamp to [0,1]
    sim = (sim + 1.0) / 2.0
    sim = np.clip(sim, 0.0, 1.0).astype(np.float32)
    return sim


def compute_similarity_sentence_transformers(texts, model_name='all-MiniLM-L6-v2'):
    try:
        from sentence_transformers import SentenceTransformer
        from sklearn.metrics.pairwise import cosine_similarity
    except Exception as e:
        raise RuntimeError('sentence-transformers or sklearn not installed') from e
    model = SentenceTransformer(model_name)
    emb = model.encode(texts, convert_to_numpy=True, show_progress_bar=False)
    sim = cosine_similarity(emb)
    sim = (sim + 1.0) / 2.0
    sim = np.clip(sim, 0.0, 1.0).astype(np.float32)
    return sim


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('--input', required=True, help='path to input JSON or JSONL file (list of samples)')
    parser.add_argument('--out', required=True, help='output directory to write train/valid/json/csv/npy')
    parser.add_argument('--split', type=float, default=0.9, help='train fraction')
    parser.add_argument('--seed', type=int, default=42)
    parser.add_argument('--sim-method', choices=['tfidf', 'st'], default='tfidf', help='similarity method: tfidf (fast) or st (sentence-transformers)')
    parser.add_argument('--st-model', default='all-MiniLM-L6-v2', help='sentence-transformers model name (if --sim-method st)')
    args = parser.parse_args()

    data = read_input(args.input)
    # Normalize entries: ensure fields name, img_path, text (text optional)
    entries = []
    for d in data:
        e = dict(d)
        if 'img_path' not in e and 'image' in e:
            e['img_path'] = e.pop('image')
        if 'name' not in e:
            e['name'] = ensure_name_field(e)
        entries.append(e)

    import random
    random.seed(args.seed)
    indices = list(range(len(entries)))
    random.shuffle(indices)
    n_train = int(len(entries) * args.split)
    train_idx = set(indices[:n_train])

    train_entries = [entries[i] for i in indices[:n_train]]
    valid_entries = [entries[i] for i in indices[n_train:]]

    out = Path(args.out)
    out.mkdir(parents=True, exist_ok=True)

    # For RadIR Stage2 we need train.jsonl/valid.jsonl entries with keys name & img_path
    train_jsonl = out / 'train.jsonl'
    valid_jsonl = out / 'valid.jsonl'
    save_jsonl([{'name':e['name'], 'img_path': e['img_path'], 'text': e.get('text','')} for e in train_entries], train_jsonl)
    save_jsonl([{'name':e['name'], 'img_path': e['img_path'], 'text': e.get('text','')} for e in valid_entries], valid_jsonl)

    # Create csv dir and npy dir
    csv_dir = out / 'csv'
    npy_dir = out / 'npy'
    csv_dir.mkdir(parents=True, exist_ok=True)
    npy_dir.mkdir(parents=True, exist_ok=True)

    # CSV should list File Path as names (use train set by default)
    names = [e['name'] for e in train_entries]
    save_csv_filepaths(names, csv_dir / 'all.csv')

    # compute similarity matrix from texts (train set rows order)
    texts = [e.get('text','') or '' for e in train_entries]
    if len(texts) == 0:
        sim = np.eye(len(names), dtype=np.float32)
    else:
        if args.sim_method == 'st':
            sim = compute_similarity_sentence_transformers(texts, model_name=args.st_model)
        else:
            sim = compute_similarity_tfidf(texts)

    np.save(npy_dir / 'all.npy', sim.astype(np.float32))

    print('Wrote:')
    print(' -', train_jsonl)
    print(' -', valid_jsonl)
    print(' -', csv_dir / 'all.csv')
    print(' -', npy_dir / 'all.npy')


if __name__ == '__main__':
    main()