"""
Inference: rank all unknown spectra by EMR probability using the 5-model ensemble.

Input : Data Parquet/unknown_data_after_noise_filter_cropped.parquet (sclk + spectrum_cropped)
Models: EMR Classifier/models/emr_model_seed_*.pth
Output: EMR Classifier/predictions/emr_candidates_unknown.csv  (all rows, sorted desc)
        EMR Classifier/predictions/emr_top100.csv              (top 100 for manual review)
"""
import os

import numpy as np
import pandas as pd
import torch
from torch.utils.data import DataLoader, TensorDataset

from train_emr_classifier import Spectra1DCNN, TARGET_LEN, minmax

ROOT = os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
HERE = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))  # EMR Classifier/
DATA_PATH = os.path.join(ROOT, 'Data Parquet', 'unknown_data_after_noise_filter_cropped.parquet')
MODELS_DIR = os.path.join(HERE, 'models')
OUT_DIR = os.path.join(HERE, 'predictions')
os.makedirs(OUT_DIR, exist_ok=True)

BATCH_SIZE = 256


def main():
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    print(f'Using device: {device}')

    print('Loading unknown data (sclk + spectrum_cropped)...')
    df = pd.read_parquet(DATA_PATH, columns=['sclk', 'spectrum_cropped'])
    X = np.stack(df['spectrum_cropped'].to_numpy()).astype(np.float32)
    X = minmax(X)
    print(f'  {X.shape}')

    model_files = sorted(f for f in os.listdir(MODELS_DIR) if f.startswith('emr_model_seed_'))
    models = []
    for mf in model_files:
        m = Spectra1DCNN(TARGET_LEN)
        m.load_state_dict(torch.load(os.path.join(MODELS_DIR, mf), map_location=device))
        m.to(device).eval()
        models.append(m)
    print(f'Loaded {len(models)} ensemble models: {model_files}')

    loader = DataLoader(TensorDataset(torch.from_numpy(X).unsqueeze(1)),
                        batch_size=BATCH_SIZE, shuffle=False, num_workers=0)

    preds = np.zeros(len(X), dtype=np.float32)
    with torch.no_grad():
        for i, (bx,) in enumerate(loader):
            bx = bx.to(device)
            acc = torch.zeros(bx.size(0), device=device)
            for m in models:
                acc += torch.sigmoid(m(bx)).squeeze(-1)
            acc /= len(models)
            s = i * BATCH_SIZE
            preds[s:s + bx.size(0)] = acc.cpu().numpy()
            if (i + 1) % 40 == 0:
                print(f'  {(i + 1) * BATCH_SIZE:,} / {len(X):,}')

    df['emr_probability'] = preds
    df_sorted = (df.drop(columns=['spectrum_cropped'])
                 .sort_values('emr_probability', ascending=False)
                 .reset_index(names='original_parquet_index'))

    out_all = os.path.join(OUT_DIR, 'emr_candidates_unknown.csv')
    df_sorted.to_csv(out_all, index=False)
    print(f'Wrote {out_all} ({len(df_sorted):,} rows)')

    out_top = os.path.join(OUT_DIR, 'emr_top100.csv')
    df_sorted.head(100).to_csv(out_top, index=False)
    print(f'Wrote {out_top}')

    print(f'\nProbability distribution: '
          f'mean={preds.mean():.4f} p50={np.percentile(preds, 50):.4f} '
          f'p99={np.percentile(preds, 99):.4f} p99.9={np.percentile(preds, 99.9):.4f} '
          f'max={preds.max():.4f}')
    for thr in (0.5, 0.75, 0.9):
        print(f'  rows with emr_probability >= {thr}: {(preds >= thr).sum():,}')


if __name__ == '__main__':
    main()