"""Train only on the provided training split. No validation labels are used."""
import argparse
import json
from pathlib import Path

import numpy as np
import pandas as pd

from model import FEATURE_NAMES, features


# def train(train_path, output):
#     frame = pd.read_parquet(train_path)
#     x = features(frame[['smiles', 'sequence']].to_dict('records'))
#     y = frame.log10_value.to_numpy(float)
#     mean, scale = x.mean(axis=0), x.std(axis=0)
#     scale[scale < 1e-10] = 1
#     x = (x - mean) / scale
#     intercept = float(y.mean())
#     coef = np.linalg.solve(x.T @ x + 10 * np.eye(x.shape[1]), x.T @ (y - intercept))
#     out = Path(output)
#     out.parent.mkdir(parents=True, exist_ok=True)
#     out.write_text(json.dumps({'feature_names': FEATURE_NAMES, 'mean': mean.tolist(),
#                               'scale': scale.tolist(), 'coef': coef.tolist(), 'intercept': intercept}, indent=2))
#     print(f'Trained baseline on {len(frame)} training rows; saved {out}')


def train(train_path, output):
    import torch  # inline so this stays a drop-in replacement

    frame = pd.read_parquet(train_path)
    x = features(frame[['smiles', 'sequence']].to_dict('records'))
    y = frame.log10_value.to_numpy(float)
    mean, scale = x.mean(axis=0), x.std(axis=0)
    scale[scale < 1e-10] = 1
    x = (x - mean) / scale
    intercept = float(y.mean())

    torch.manual_seed(0)
    xt = torch.tensor(x, dtype=torch.float32)
    yt = torch.tensor(y - intercept, dtype=torch.float32)[:, None]
    w = torch.zeros(x.shape[1], 1, requires_grad=True)
    opt = torch.optim.AdamW([w], lr=1e-2, weight_decay=1e-3)
    sched = torch.optim.lr_scheduler.CosineAnnealingLR(opt, T_max=500)
    for _ in range(500):
        perm = torch.randperm(len(xt))
        for i in range(0, len(xt), 1024):
            b = perm[i:i + 1024]
            opt.zero_grad()
            torch.nn.functional.huber_loss(xt[b] @ w, yt[b]).backward()
            opt.step()
        sched.step()
    coef = w.detach().numpy().ravel().astype(float)

    out = Path(output)
    out.parent.mkdir(parents=True, exist_ok=True)
    out.write_text(json.dumps({'feature_names': FEATURE_NAMES, 'mean': mean.tolist(),
                              'scale': scale.tolist(), 'coef': coef.tolist(), 'intercept': intercept}, indent=2))
    print(f'Trained baseline on {len(frame)} training rows; saved {out}')


if __name__ == '__main__':
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument('--train', default='../data/public/train.parquet')
    parser.add_argument('--output', default='weights/baseline.json')
    args = parser.parse_args()
    train(args.train, args.output)