import argparse
import json

import numpy as np
import pandas as pd

from inputs import prepare_inputs

parser = argparse.ArgumentParser()
parser.add_argument('--t', type=float, default=0.05)
parser.add_argument('--frac', type=float, default=0.1)
parser.add_argument('--seed', type=int, default=0)
parser.add_argument('--output', default='hard_split.json')
args = parser.parse_args()

K, MAXDF = 5, 2000
LUT = np.full(256, 20, dtype=np.int64)
for i, c in enumerate('ACDEFGHIKLMNPQRSTVWY'):
    LUT[ord(c)] = i


def kmers(seq):
    codes = LUT[np.frombuffer(seq.encode('ascii', 'ignore'), dtype=np.uint8)]
    n = len(codes) - K + 1
    if n <= 0:
        return np.empty(0, dtype=np.int32)
    ids = np.zeros(n, dtype=np.int64)
    for j in range(K):
        ids = ids * 21 + codes[j:j + n]
    return np.unique(ids).astype(np.int32)


train = pd.read_parquet('../data/train.parquet').drop_duplicates(['sequence', 'ec_number'])
seqs = train.groupby('sequence', sort=False).size().index.tolist()
N = len(seqs)
parts = [kmers(s) for s in seqs]
sizes = np.array([len(p) for p in parts])
kmer_all = np.concatenate(parts)
del parts
seq_all = np.repeat(np.arange(N, dtype=np.int32), sizes)
order = np.argsort(kmer_all, kind='stable')
kmer_sorted = kmer_all[order]
seq_sorted = seq_all[order]
del kmer_all, seq_all, order
print('index built for', N, 'sequences')


def shared_counts(seq):
    q = kmers(seq)
    if len(q) == 0:
        return np.zeros(N, dtype=np.int64), 1
    lo = np.searchsorted(kmer_sorted, q, 'left')
    cnt = np.searchsorted(kmer_sorted, q, 'right') - lo
    keep = cnt <= MAXDF
    lo, cnt = lo[keep], cnt[keep]
    offs = np.cumsum(cnt) - cnt
    idx = np.arange(cnt.sum()) - np.repeat(offs, cnt) + np.repeat(lo, cnt)
    return np.bincount(seq_sorted[idx], minlength=N), len(q)


def nearest(seq, blocked=None):
    shared, nq = shared_counts(seq)
    if blocked is not None:
        shared[blocked] = 0
    return shared.max() / nq


rng = np.random.RandomState(args.seed)
target = int(args.frac * N)
in_hold = np.zeros(N, dtype=bool)
for s in rng.permutation(N):
    if in_hold[s]:
        continue
    shared, nq = shared_counts(seqs[s])
    group = (shared >= args.t * nq) & ~in_hold
    group[s] = True
    if in_hold.sum() + group.sum() > 1.2 * target:
        continue
    in_hold |= group
    if in_hold.sum() >= target:
        break
print('holdout sequences:', int(in_hold.sum()), 'of', N)

val_seqs = prepare_inputs(pd.read_parquet('../data/val_inputs.parquet')).sequence.tolist()
hold_ids = np.where(in_hold)[0]
sample = rng.choice(hold_ids, min(3000, len(hold_ids)), replace=False)
random_ids = rng.choice(N, 3000, replace=False)

rows = []
for name, sims in (
        ('validation to full train', [nearest(s) for s in val_seqs]),
        ('new holdout to remaining train', [nearest(seqs[i], in_hold) for i in sample]),
        ('random holdout to remaining train', [nearest(seqs[i], np.arange(N) == i) for i in random_ids])):
    p = np.percentile(sims, [10, 25, 50, 75, 90])
    rows.append({'set': name, 'p10': p[0], 'p25': p[1], 'p50': p[2], 'p75': p[3], 'p90': p[4]})
print(pd.DataFrame(rows).round(3).to_string(index=False))

json.dump({'t': args.t, 'holdout': [int(i) for i in hold_ids]}, open(args.output, 'w'))
print('saved', args.output)