import argparse
import json
import time
from pathlib import Path

import numpy as np
import pandas as pd
import torch

import dlmodel as dl

parser = argparse.ArgumentParser()
parser.add_argument('--train', required=True)
parser.add_argument('--output', default='weights')
parser.add_argument('--epochs', type=int, default=12)
parser.add_argument('--holdout', type=float, default=0.1)
parser.add_argument('--split', default='')
parser.add_argument('--alpha', type=float, default=0.7)
parser.add_argument('--drop', type=float, default=0.3)
parser.add_argument('--wd', type=float, default=0.05)
parser.add_argument('--noise', type=float, default=0.1)
parser.add_argument('--crop', type=float, default=0.6)
args = parser.parse_args()

BATCH, LR, CLIP = 64, 2e-3, 1.0
dl.NOISE_P = args.noise
dl.MINKEEP = args.crop

dl.seed_all()
dev = dl.device
cfg = dl.CFG
print('device:', dev)

df = pd.read_parquet(args.train).drop_duplicates(['sequence', 'ec_number'])
vocab_file = Path(__file__).parent / 'ec_vocabulary.json'
if vocab_file.exists():
    assert set(df.ec_number) <= set(json.loads(vocab_file.read_text())), 'EC outside the vocabulary'
grouped = df.groupby('sequence', sort=False)['ec_number'].apply(lambda s: sorted(set(s)))
all_seqs, all_ecs = grouped.index.tolist(), grouped.tolist()
ec_list = sorted(df.ec_number.unique())
ec_index = {e: i for i, e in enumerate(ec_list)}
n_classes = len(ec_list)

if args.split:
    hold_ids = np.random.RandomState(0).permutation(np.array(json.load(open(args.split))['holdout']))
    train_ids = np.setdiff1d(np.arange(len(all_seqs)), hold_ids)
else:
    perm = np.random.RandomState(0).permutation(len(all_seqs))
    n = int(len(all_seqs) * args.holdout)
    hold_ids, train_ids = perm[:n], perm[n:]
n_hold = len(hold_ids)
tr_seqs = [all_seqs[i] for i in train_ids]
tr_labels = [[ec_index[e] for e in all_ecs[i]] for i in train_ids]
tr_lens = [len(s) for s in tr_seqs]
hold_seqs = [all_seqs[i] for i in hold_ids]
hold_ecs = [all_ecs[i] for i in hold_ids]
print(f'{len(tr_seqs)} training sequences, {n_hold} held out, {n_classes} classes')


def make_targets(batch_ids):
    rows, cols, vals = [], [], []
    for r, i in enumerate(batch_ids):
        labs = tr_labels[i]
        rows += [r] * len(labs)
        cols += labs
        vals += [1.0 / len(labs)] * len(labs)
    t = torch.zeros(len(batch_ids), n_classes, device=dev)
    t[torch.tensor(rows, device=dev), torch.tensor(cols, device=dev)] = torch.tensor(vals, device=dev)
    return t


def holdout_l4(P, n=5000):
    if not n_hold:
        return float('nan')
    _, probs = dl.run(P, cfg, hold_seqs[:n], dev)
    preds = [ec_list[i] for i in probs.argmax(axis=1)]
    return dl.level_scores(preds, hold_ecs[:n])[4]


P = dl.init_params(cfg, n_classes, dev)
S = dl.adam_init(P)
print('parameters:', sum(v.numel() for v in P.values()))

total_steps = len(dl.make_batches(tr_lens, BATCH, True)) * args.epochs
step = 0
best_val, best_epoch, best_P = -1.0, 0, None

for epoch in range(1, args.epochs + 1):
    t0, total, count = time.time(), 0.0, 0
    for b in dl.make_batches(tr_lens, BATCH, True):
        x = dl.make_batch([tr_seqs[i] for i in b], True, dev)
        logits, _, C = dl.net_fwd(x, P, cfg, args.drop)
        loss, dlogits = dl.cross_entropy_soft_loss_and_grad(logits, make_targets(b))
        G = dl.net_bwd(dlogits, P, cfg, C)
        norm = dl.adam_step(P, G, S, dl.lr_at(step, total_steps, LR), CLIP, args.wd)
        if not np.isfinite(norm):
            raise SystemExit('gradient is nan, try a smaller LR in train.py')
        step += 1
        total += loss.item() * len(b)
        count += len(b)
        if step % 500 == 0:
            print(f'  step {step}/{total_steps}  loss {total / count:.4f}  {time.time() - t0:.0f}s', flush=True)
    val = holdout_l4(P)
    if n_hold and val > best_val:
        best_val, best_epoch = val, epoch
        best_P = {k: v.clone() for k, v in P.items()}
    print(f'epoch {epoch}/{args.epochs}  loss {total / count:.4f}  holdout L4 {val:.4f}  {time.time() - t0:.0f}s', flush=True)

if best_P is not None:
    P = best_P
    print(f'using epoch {best_epoch} with holdout L4 {best_val:.4f}')

bank_emb, _ = dl.run(P, cfg, tr_seqs, dev, bs=128, want_probs=False)
bank = dl.make_bank(bank_emb, tr_labels, ec_list)

best_alpha = args.alpha
if n_hold:
    Pc = {k: v.cpu() for k, v in P.items()}
    torch.set_num_threads(2)
    t0 = time.time()
    emb, probs = dl.run(Pc, cfg, hold_seqs, 'cpu')
    votes = dl.knn_votes(bank, emb)
    print(f'CPU inference on {len(hold_seqs)} held out sequences: {time.time() - t0:.0f}s')
    best = -1.0
    for alpha in (0.0, 0.3, 0.5, 0.7, 1.0):
        sc = dl.level_scores(dl.pick(probs, votes, ec_list, alpha), hold_ecs)
        print(f'alpha={alpha}  ' + '  '.join(f'L{l}={sc[l]:.4f}' for l in (4, 3, 2, 1)))
        if sc[4] > best:
            best, best_alpha = sc[4], alpha
    print('best alpha:', best_alpha)

dl.save(P, cfg, bank_emb, tr_labels, ec_list, args.output, best_alpha)
print('saved to', args.output)