"""Small CPU baseline with JSON weights. Replace with your own model implementation."""
import json
from pathlib import Path
import re

import numpy as np

AMINO_ACIDS = 'ACDEFGHIKLMNPQRSTVWY'
SMILES_CHARS = 'BCNOFPSIcnops[]=#()+-@/\\1234567890'
FEATURE_NAMES = ['log_sequence_length', 'log_smiles_length'] + list(AMINO_ACIDS) + list(SMILES_CHARS)


def features(rows):
    result = []
    for row in rows:
        seq, smi = row['sequence'], row['smiles']
        seq_upper = seq.upper()
        result.append([np.log1p(len(seq)), np.log1p(len(smi))]
                      + [seq_upper.count(a) / max(1, len(seq)) for a in AMINO_ACIDS]
                      + [smi.count(a) / max(1, len(smi)) for a in SMILES_CHARS])
    return np.asarray(result, dtype=float)


class Model:
    def __init__(self, weights=None):
        path = Path(weights) if weights else Path(__file__).parent / 'weights' / 'baseline.json'
        self.weights = json.loads(path.read_text())

    def predict(self, rows):
        x = features(rows)
        w = self.weights
        return (((x - np.asarray(w['mean'])) / np.asarray(w['scale'])) @ np.asarray(w['coef'])
                + w['intercept']).tolist()


def validate_rows(rows, max_rows=256):
    if not isinstance(rows, list) or not 1 <= len(rows) <= max_rows:
        raise ValueError(f'Provide between 1 and {max_rows} rows.')
    ids = set()
    for row in rows:
        if not isinstance(row, dict) or set(row) != {'id', 'smiles', 'sequence'}:
            raise ValueError('Each row must contain exactly id, smiles, sequence.')
        if not isinstance(row['id'], str) or not re.fullmatch(r'[A-Za-z0-9_-]{1,80}', row['id']):
            raise ValueError('Invalid row ID.')
        if row['id'] in ids:
            raise ValueError('Duplicate row ID.')
        ids.add(row['id'])
        for key, limit in [('sequence', 50000), ('smiles', 20000)]:
            if not isinstance(row[key], str) or not 1 <= len(row[key]) <= limit:
                raise ValueError(f'Invalid {key}.')
    return rows