"""LightGBM on RDKit substrate features, sequence composition, and ESM-2 (8M) embeddings."""
import json
from pathlib import Path
import re

import numpy as np

from featurize import composition_features, esm_embeddings, smiles_features

WEIGHTS = Path(__file__).parent / 'weights'


class Model:
    def __init__(self, weights=WEIGHTS, threads=2):
        import lightgbm as lgb

        self.meta = json.loads((Path(weights) / 'meta.json').read_text())
        self.threads = threads
        self.boosters = [lgb.Booster(model_file=str(Path(weights) / f'lgbm_{t}.txt')) for t in self.meta['targets']]

    def predict(self, frame):
        """Return an (n, 2) array of log10_kcat, log10_km for a frame with smiles and sequence."""
        blocks = []
        groups = self.meta['groups']
        if 'smiles' in groups:
            blocks.append(self._per_unique(frame.smiles, smiles_features))
        if 'comp' in groups:
            blocks.append(self._per_unique(frame.sequence, composition_features))
        if 'esm' in groups:
            blocks.append(self._per_unique(frame.sequence, lambda s: esm_embeddings(s, threads=self.threads)))
        x = np.hstack(blocks)
        return np.column_stack([b.predict(x, num_threads=self.threads) for b in self.boosters])

    @staticmethod
    def _per_unique(values, fn):
        keys, inverse = np.unique(values.to_numpy(), return_inverse=True)
        return fn(list(keys))[inverse]


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