#!/usr/bin/env python3
"""Reproduce Veritiana AI Meter classifier v3.1.0 from training-balanced.jsonl."""
from __future__ import annotations
import argparse, hashlib, json, struct, sys, time
from collections import Counter
from pathlib import Path
import numpy as np
from sklearn.linear_model import LogisticRegression
from sklearn.metrics import classification_report, confusion_matrix
from sklearn.model_selection import GroupShuffleSplit
from features import INPUT_SIZE, build_features

TASK_LABELS = ["general_chat","writing","translation","summarization","research","coding","mathematics","document_analysis","high_stakes"]
COMPLEXITY_LABELS = ["low","medium","high"]

def sha256(path: Path) -> str:
    return hashlib.sha256(path.read_bytes()).hexdigest()

def _varint(number: int) -> bytes:
    number = int(number)
    if number < 0: number = (1 << 64) + number
    output = bytearray()
    while True:
        byte = number & 0x7F; number >>= 7
        output.append(byte | (0x80 if number else 0))
        if not number: return bytes(output)
def _key(field: int, wire: int) -> bytes: return _varint((field << 3) | wire)
def _fv(field: int, value: int) -> bytes: return _key(field,0)+_varint(value)
def _fb(field: int, value: bytes) -> bytes: return _key(field,2)+_varint(len(value))+value
def _fs(field: int, value: str) -> bytes: return _fb(field,value.encode())
def _msg(field: int, value: bytes) -> bytes: return _fb(field,value)
def _packed(field: int, values) -> bytes: return _fb(field,b''.join(_varint(v) for v in values))
def _tensor(name: str, array: np.ndarray) -> bytes:
    array=np.ascontiguousarray(array.astype(np.float32))
    return _packed(1,array.shape)+_fv(2,1)+_fs(8,name)+_fb(9,array.tobytes(order='C'))
def _dim(value): return _fs(2,value) if isinstance(value,str) else _fv(1,value)
def _shape(dims): return b''.join(_msg(1,_dim(value)) for value in dims)
def _tensor_type(element,dims): return _fv(1,element)+_msg(2,_shape(dims))
def _type_proto(element,dims): return _msg(1,_tensor_type(element,dims))
def _value_info(name,dims): return _fs(1,name)+_msg(2,_type_proto(1,dims))
def _attr_int(name,value): return _fs(1,name)+_fv(3,value)+_fv(20,2)
def _node(inputs,outputs,operation,name='',attributes=None):
    value=b''.join(_fs(1,item) for item in inputs)+b''.join(_fs(2,item) for item in outputs)
    if name: value += _fs(3,name)
    value += _fs(4,operation)
    for attribute in attributes or []: value += _msg(5,attribute)
    return value
def _graph(nodes,initializers,inputs,outputs,name):
    return (b''.join(_msg(1,node) for node in nodes)+_fs(2,name)+
            b''.join(_msg(5,item) for item in initializers)+
            b''.join(_msg(11,item) for item in inputs)+
            b''.join(_msg(12,item) for item in outputs))
def _opset(version): return _fv(2,version)
def _kv(key,value): return _fs(1,key)+_fs(2,value)

def export_onnx(task_model, complexity_model, destination: Path, version: str) -> None:
    initializers=[
        _tensor('task_weights',task_model.coef_.astype(np.float32).T.copy()),
        _tensor('task_bias',task_model.intercept_.astype(np.float32).copy()),
        _tensor('complexity_weights',complexity_model.coef_.astype(np.float32).T.copy()),
        _tensor('complexity_bias',complexity_model.intercept_.astype(np.float32).copy()),
    ]
    nodes=[
        _node(['features','task_weights'],['task_mm'],'MatMul','TaskMatMul'),
        _node(['task_mm','task_bias'],['task_logits'],'Add','TaskAdd'),
        _node(['task_logits'],['task_probabilities'],'Softmax','TaskSoftmax',[_attr_int('axis',1)]),
        _node(['features','complexity_weights'],['complexity_mm'],'MatMul','ComplexityMatMul'),
        _node(['complexity_mm','complexity_bias'],['complexity_logits'],'Add','ComplexityAdd'),
        _node(['complexity_logits'],['complexity_probabilities'],'Softmax','ComplexitySoftmax',[_attr_int('axis',1)]),
    ]
    graph=_graph(nodes,initializers,[_value_info('features',['batch',1544])],
                 [_value_info('task_probabilities',['batch',9]),_value_info('complexity_probabilities',['batch',3])],
                 'VeritianaAIMeterClassifier')
    metadata={
        'model':'Veritiana Multilingual Intent Classifier','version':version,'input_size':'1544',
        'task_labels':json.dumps(TASK_LABELS),'complexity_labels':json.dumps(COMPLEXITY_LABELS),
        'training_data':'OASST1 + CoEdIT + MBPP+ + IFEval; filtered, weak-labeled, balanced',
        'feature_contract':'AI Meter v1.4.9 compatible'
    }
    payload=(_fv(1,8)+_fs(2,'veritiana-ai-meter-model-trainer')+_fs(3,version)+
             _msg(7,graph)+_msg(8,_opset(13))+b''.join(_msg(14,_kv(k,v)) for k,v in metadata.items()))
    destination.write_bytes(payload)

def main() -> int:
    parser=argparse.ArgumentParser()
    parser.add_argument('dataset',type=Path,help='Prepared training-balanced.jsonl')
    parser.add_argument('--output-dir',type=Path,default=Path('output'))
    parser.add_argument('--version',default='3.1.0-multisource-balanced')
    parser.add_argument('--test-size',type=float,default=0.20)
    parser.add_argument('--seed',type=int,default=42)
    parser.add_argument('--c',type=float,default=4.0)
    args=parser.parse_args()
    rows=[json.loads(line) for line in args.dataset.read_text(encoding='utf-8').splitlines() if line.strip()]
    if not rows: raise ValueError('Empty dataset')
    task_map={label:index for index,label in enumerate(TASK_LABELS)}
    complexity_map={label:index for index,label in enumerate(COMPLEXITY_LABELS)}
    X=np.vstack([build_features(row['text']) for row in rows]).astype(np.float32)
    task_y=np.asarray([task_map[row['task']] for row in rows],dtype=np.int64)
    complexity_y=np.asarray([complexity_map[row['complexity']] for row in rows],dtype=np.int64)
    groups=[row.get('derived_from') or row.get('normalized_hash') or row.get('id') for row in rows]
    splitter=GroupShuffleSplit(n_splits=1,test_size=args.test_size,random_state=args.seed)
    train_index,test_index=next(splitter.split(X,task_y,groups=groups))
    def fit(target):
        model=LogisticRegression(C=args.c,max_iter=3000,solver='lbfgs',class_weight='balanced',random_state=args.seed)
        model.fit(X[train_index],target[train_index]); return model
    started=time.perf_counter(); task_model=fit(task_y); complexity_model=fit(complexity_y); seconds=time.perf_counter()-started
    task_prediction=task_model.predict(X[test_index]); complexity_prediction=complexity_model.predict(X[test_index])
    task_report=classification_report(task_y[test_index],task_prediction,labels=list(range(9)),target_names=TASK_LABELS,output_dict=True,zero_division=0)
    complexity_report=classification_report(complexity_y[test_index],complexity_prediction,labels=list(range(3)),target_names=COMPLEXITY_LABELS,output_dict=True,zero_division=0)
    evaluation={
        'summary':{'task_accuracy':task_report['accuracy'],'task_macro_f1':task_report['macro avg']['f1-score'],
                   'complexity_accuracy':complexity_report['accuracy'],'complexity_macro_f1':complexity_report['macro avg']['f1-score'],
                   'training_seconds':seconds},
        'task_report':task_report,'complexity_report':complexity_report,
        'task_confusion_matrix':confusion_matrix(task_y[test_index],task_prediction,labels=list(range(9))).tolist(),
        'complexity_confusion_matrix':confusion_matrix(complexity_y[test_index],complexity_prediction,labels=list(range(3))).tolist()
    }
    args.output_dir.mkdir(parents=True,exist_ok=True)
    model_path=args.output_dir/'veritiana-classifier-v2.onnx'; export_onnx(task_model,complexity_model,model_path,args.version)
    (args.output_dir/'evaluation.json').write_text(json.dumps(evaluation,ensure_ascii=False,indent=2),encoding='utf-8')
    meta={
        'model':'Veritiana Multilingual Intent Classifier','version':args.version,'engine':'ONNX Runtime Web',
        'architecture':'1544 hashed lexical/character/numeric features; dual multinomial logistic heads',
        'input':{'name':'features','dtype':'float32','shape':['batch',1544]},
        'outputs':[{'name':'task_probabilities','dtype':'float32','shape':['batch',9]},
                   {'name':'complexity_probabilities','dtype':'float32','shape':['batch',3]}],
        'task_labels':TASK_LABELS,'complexity_labels':COMPLEXITY_LABELS,
        'training_data':{'dataset_sha256':sha256(args.dataset),'balanced_rows':len(rows),'train_rows':len(train_index),'test_rows':len(test_index),
                         'label_method_counts':dict(Counter(row.get('label_method','unknown') for row in rows))},
        'validation':evaluation['summary'],'onnx_sha256':sha256(model_path),
        'limitations':'Validation uses weak labels and deterministic augmentations; it is not an independent human-ground-truth benchmark.'
    }
    (args.output_dir/'classifier-meta.json').write_text(json.dumps(meta,ensure_ascii=False,indent=2),encoding='utf-8')
    print(json.dumps({'model':str(model_path),'onnx_sha256':sha256(model_path),'evaluation':evaluation['summary']},indent=2))
    return 0
if __name__=='__main__': raise SystemExit(main())