"""Build a standalone illustrated LOO report from saved predictions and real inputs."""
from __future__ import annotations

import argparse
import base64
from datetime import datetime, timezone
import hashlib
import json
import os
from pathlib import Path
import sys
import time

if __package__ in (None, ''):
    sys.path.insert(0, str(Path(__file__).resolve().parents[1]))

import numpy as np
from common import ROOT, read


def relative(path):
    return str(Path(path).resolve().relative_to(ROOT))


def atomic(path, content):
    path = Path(path); path.parent.mkdir(parents=True, exist_ok=True)
    temporary = path.with_suffix(path.suffix + '.tmp')
    temporary.write_text(content, encoding='utf-8'); temporary.replace(path)


def json_text(value):
    return json.dumps(value, ensure_ascii=False, separators=(',', ':'), allow_nan=False)


def condition(record):
    """Same fixed scales and validity channels used by Models.adapters.Sample."""
    with np.load(record['process']) as data:
        names = data['feature_names'].astype(str).tolist()
        raw = data['x'].copy(); valid = data['feature_valid'].astype(bool)
    values = raw.copy()
    for i, name in enumerate(names):
        unit = name.rsplit('[', 1)[-1].rstrip(']').lower()
        if unit in ('°c', 'c', 'degc'): values[i] /= 500.
        elif unit in ('s', 'sec'): values[i] /= 100.
        else: values[i] = np.sign(values[i]) * np.log1p(abs(values[i]))
    return {'names': names, 'raw': raw.tolist(), 'scaled': values.tolist(), 'valid': valid.tolist(),
            'shape': [2 * len(names)], 'source': relative(record['process'])}


def extract_case(record, profile, source):
    """Only eligible unseen-design tests; never substitute native metrics for null reference metrics."""
    if record['protocol'] != 'leave_one_design_out' or record.get('training_excluded') or not record.get('trained'):
        return None
    fit = record.get('common_grid_weak_fit') if 'common_grid_weak_fit' in record else record.get('weak_fit')
    if not fit: return None
    grid = record.get('common_grid') or record
    y = grid.get('target')
    prediction = grid.get('predictions')
    mask = [bool(a and b and v is not None) for a, b, v in zip(grid['training_valid'], grid['prediction_valid'], y)]
    if not (len(y) == len(prediction) == len(mask)): raise ValueError('Evaluation grid dimensions differ')
    return {'id': f"{profile}/{record['model_id']}/{record['dataset_id']}", 'profile': profile,
            'model_id': record['model_id'], 'dataset_id': record['dataset_id'], 'design_id': record['design_id'],
            'fold_id': record['fold_id'], 'production_status': record['production_status'],
            'input_size': record['input_size'], 'actual_input_shape': record.get('actual_input_shape'),
            'native_nodes': record.get('actual_cad_nodes'), 'metrics': fit, 'predictions': prediction,
            'target': y, 'valid': mask, 'threshold': record['config'].get('score_threshold', .5),
            'source': relative(source), 'prediction_source': relative(record['model_prediction_source']),
            'run': record['date'] + '/' + record['run_id'], 'label_status': record['label_status'],
            'train_count': len(record.get('actual_train_dataset_ids', [])),
            'grid': 'Frozen reference CAD nodes' if record.get('common_grid') else 'Native CAD nodes'}


def stats(values):
    valid = [float(v) for v in values if v is not None and np.isfinite(v)]
    return {'n': len(valid), 'mean': float(np.mean(valid)) if valid else None,
            'min': min(valid) if valid else None, 'max': max(valid) if valid else None}


def summarize(cases):
    groups = {}
    for case in cases: groups.setdefault((case['profile'], case['model_id']), []).append(case)
    return [{'profile': profile, 'model_id': model, 'cases': len(group),
             **{key: stats([c['metrics'].get(key) for c in group]) for key in ('bce', 'average_precision', 'iou')}}
            for (profile, model), group in sorted(groups.items())]


def preview(design, model, config, cache_dir):
    """Scientific previews of saved arrays, with explicit display sampling and actual dimensions."""
    from Models.registry import BY_ID
    kind = BY_ID[model]['representation']
    key = f"img_{config['image_sizes'][0]}" if kind == 'img' else f"voxel_{config['voxel_sizes'][0]}" if kind == 'voxel' else 'brep' if kind == 'brep' else 'graph'
    source = Path(design['representations'][key]); stamp = source.stat()
    digest = hashlib.sha256(f'{source}:{stamp.st_size}:{stamp.st_mtime_ns}:{kind}:v1'.encode()).hexdigest()[:20]
    output = cache_dir / (digest + '.png'); metadata = output.with_suffix('.json')
    if output.is_file() and metadata.is_file():
        result = read(metadata)
        return {**result, 'png': 'data:image/png;base64,' + base64.b64encode(output.read_bytes()).decode()}
    import matplotlib
    matplotlib.use('Agg')
    import matplotlib.pyplot as plt
    plt.rcParams.update({'font.size': 9, 'axes.titlesize': 10})
    description = ''; shape = None
    with np.load(source) as data:
        if kind == 'img':
            size = config['image_sizes'][0]; stride = max(1, int(np.ceil(size/96)))
            depth = data['depth_mm'][:, ::stride, ::stride] / 1000.
            valid = data['feature_valid'][:, 3, ::stride, ::stride]
            shape = [6, 34, size, size]
            fig, axes = plt.subplots(2, 3, figsize=(8, 5))
            values = depth[valid]
            lo, hi = (float(values.min()), float(values.max())) if len(values) else (0, 1)
            for i, ax in enumerate(axes.flat):
                ax.imshow(np.ma.masked_where(~valid[i], depth[i]), cmap='viridis', vmin=lo, vmax=hi)
                ax.set_title(['+X', '-X', '+Y', '-Y', '+Z', '-Z'][i]); ax.axis('off')
            description = f'Six actual CAD raster views, depth channel /1000 (channel 3 of 34); invalid pixels blank. Display stride {stride}; model uses full {size}×{size}. The separate process vector also enters CNN/FNO.'
        elif kind == 'voxel':
            x = data['x']; volume = x[0]; size = config['voxel_sizes'][0]; shape = [1, 4, size, size, size]
            fig, axes = plt.subplots(1, 3, figsize=(8, 3))
            for axis, ax in enumerate(axes):
                ax.imshow(volume.max(axis=axis), cmap='Blues', vmin=0, vmax=1, origin='lower')
                ax.set_title('Occupancy max projection ' + ['Y/Z', 'X/Z', 'X/Y'][axis]); ax.axis('off')
            description = 'Exact saved sampled-surface occupancy channel (one of four voxel channels), maximum projected along each axis. This is surface voxelization, not filled solid occupancy. Process conditions enter separately.'
        elif kind == 'brep':
            trim = data['trim_valid']; face_ids = np.unique(data['node_face'])
            selected = face_ids[np.linspace(0, len(face_ids)-1, min(6, len(face_ids)), dtype=int)]
            pos = data['pos_mm']; face_z = pos[selected, :, :, 2]/1000.
            shape = [trim.shape[0], 8 + config['process_dim'], trim.shape[1], trim.shape[2]]
            fig, axes = plt.subplots(2, 3, figsize=(8, 5))
            for ax in axes.flat: ax.axis('off')
            for i, fid in enumerate(selected):
                axes.flat[i].imshow(np.ma.masked_where(~trim[fid], face_z[i]), cmap='viridis')
                axes.flat[i].set_title(f'CAD face {fid} · Z/1000')
            description = f'Six representative actual trimmed B-Rep UV faces, Z/1000 channel; each face has its own color scale. {trim.shape[0]} faces enter the model at {trim.shape[1]}×{trim.shape[2]}, with 8 geometry/mask channels plus {config["process_dim"]} process/mask channels. Face adjacency is also input.'
        else:
            pos = data['pos_mm']; n = len(pos); shape = [n, 12 if kind == 'graph' else 12+config['process_dim']]
            fig, axes = plt.subplots(1, 3, figsize=(8, 3))
            edges = data['edge_index'].T if kind == 'graph' else []
            if len(edges) > 8000: edges = edges[np.linspace(0, len(edges)-1, 8000, dtype=int)]
            for ax, pair, label in zip(axes, [(0,2), (0,1), (1,2)], ['X/Z', 'X/Y', 'Y/Z']):
                if len(edges):
                    from matplotlib.collections import LineCollection
                    ax.add_collection(LineCollection(pos[edges][:,:,pair], colors='#c5d1d9', linewidths=.25))
                ax.scatter(pos[:,pair[0]], pos[:,pair[1]], c=pos[:,2]/1000, s=3, cmap='viridis')
                ax.set_aspect('equal', adjustable='datalim'); ax.set_title(label+' · mm'); ax.tick_params(labelsize=7)
            description = f'All {n} actual CAD input nodes/points, colored by Z/1000. '
            if kind == 'graph': description += f'Edges shown: {len(edges)} (display cap 8000). Input has 12 node features, 4 edge features, and {config["process_dim"]} separate process/mask values.'
            else: description += f'Each node receives 12 geometry/mask features plus {config["process_dim"]} process/mask values ({12+config["process_dim"]} columns total).'
            if kind == 'tabular': description += 'These are the spatial rows of the tabular feature matrix; exact process values appear in the inspector.'
    fig.tight_layout(); output.parent.mkdir(parents=True, exist_ok=True); fig.savefig(output, dpi=125, facecolor='white'); plt.close(fig)
    result = {'source': relative(source), 'kind': kind, 'description': description, 'shape': shape}
    atomic(metadata, json_text(result))
    return {**result, 'png': 'data:image/png;base64,'+base64.b64encode(output.read_bytes()).decode()}


def build(series_path):
    from Models.registry import BY_ID
    series = read(series_path); directory = ROOT/'reports'/Path(series['directory']).parent.name/Path(series['directory']).name
    cache_dir = directory/'input_previews'
    reference = read(series['reference_prepared_manifest'])
    geometry = {}
    for did, design in reference['designs'].items():
        with np.load(design['representations']['graph']) as graph:
            geometry[did] = {'positions': graph['pos_mm'].tolist(), 'node_ids': graph['node_id'].tolist()}
    cases, inputs, conditions, statuses, profiles = [], {}, {}, [], []
    for level in series['levels']:
        profile = level['image_size']; profiles.append({k:level.get(k) for k in ['image_size','node_budget','voxel_size','uv_grid','gp_max_train_nodes','phase']})
        if not level.get('run'): continue
        prepared = read(Path(level['run']['manifest_dir'])/'prepared_manifest.json')
        records = {r['dataset_id']:r for r in prepared['datasets']}
        config = {**level['run']['config'], 'process_dim': 2*len(prepared['process_schema'])}
        for model, result in level.get('model_results', {}).items():
            statuses.append({'profile':profile,'model_id':model,'status_counts':result['status_counts'], 'input_size':result['input_size']})
            for source in result['source_metric_files']:
                if '/leave_one_design_out/' not in source: continue
                record = read(source); case = extract_case(record, profile, source)
                if case is None: continue
                dataset = records[case['dataset_id']]; design = prepared['designs'][case['design_id']]
                key = f"{profile}/{case['design_id']}/{BY_ID[model]['representation']}"
                if key not in inputs:
                    inputs[key] = preview(design, model, config, cache_dir)
                case['input_key'] = key
                if case['dataset_id'] not in conditions: conditions[case['dataset_id']] = condition(dataset)
                cases.append(case)
    payload = {'generated_at':datetime.now(timezone.utc).isoformat(), 'epochs':series['epochs'], 'models':series['models'],
               'profiles':profiles,'cases':cases,'summary':summarize(cases),'inputs':inputs,'geometry':geometry,
               'conditions':conditions,'statuses':statuses,'series_status':series['status'],
               'exclusions':reference['run']['config'].get('excluded_train_dataset_ids',[]),
               'source':relative(series_path), 'reference':relative(series['reference_prepared_manifest']),
               'process_dim':2*len(reference['process_schema'])}
    template = (ROOT/'Summary/report_template.html').read_text()
    data = json_text(payload).replace('<','\\u003c')
    # Scripts/styles/data are embedded so this HTML also opens offline.
    content = template.replace('/* REPORT_CSS */',(ROOT/'Summary/report.css').read_text()).replace('/* REPORT_JS */',(ROOT/'Summary/report.mjs').read_text().replace('export ',''))
    content = content.replace('REPORT_DATA_JSON',data)
    atomic(directory/'results.html',content)
    atomic(directory/'report_status.json',json_text({'generated_at':payload['generated_at'],'cases':len(cases)}))
    atomic(ROOT/'reports/loo_report.html','<!doctype html><meta http-equiv="refresh" content="0;url='+str((directory/'results.html').relative_to(ROOT/'reports'))+'"><a href="'+str((directory/'results.html').relative_to(ROOT/'reports'))+'">Open illustrated LOO report</a>')
    for level in series['levels']:
        run = level.get('run')
        if not run: continue
        for path in [ROOT/'reports'/run['date']/run['run_id']/'results.html', ROOT/'Summary/outputs'/run['date']/run['run_id']/'index.html']:
            if not path.is_file(): continue
            old = path.read_text(); marker = 'id="illustrated-loo-report"'
            if marker in old: continue
            url = os.path.relpath(directory/'results.html',path.parent)
            link = '<p id="illustrated-loo-report"><a href="'+url+'">Illustrated LOO report: input sizes, scores, best/worst cases and metric explanations ↗</a></p>'
            index = old.find('<h1>')
            atomic(path,old[:index]+link+old[index:] if index >= 0 else old+link)
    atomic(directory/'build.json',json_text({'generated_at':payload['generated_at'],'cases':len(cases),'model_sizes':len(payload['summary']),
        'input_previews':len(inputs),'profiles':[(l['image_size'],l['phase']) for l in series['levels']]}))
    print(f'Built {directory}/results.html · {len(cases)} LOO cases · {len(inputs)} actual-input previews',flush=True)
    return directory


def main():
    parser = argparse.ArgumentParser(description=__doc__); parser.add_argument('--series',type=Path); parser.add_argument('--watch',action='store_true')
    args = parser.parse_args(); series_path = args.series or Path(read(ROOT/'Experiments/latest.json')['series'])
    previous = None
    while True:
        series = None
        try:
            series = read(series_path)
            source_stamps = [(name, (ROOT/'Summary'/name).stat().st_mtime_ns) for name in ('report_template.html','report.mjs','report.css','visual_report.py')]
            signature = json_text([source_stamps, [(l['image_size'],l['phase'],list(l.get('model_results',{}))) for l in series['levels']]])
            if signature != previous:
                build(series_path); previous = signature
        except Exception as exc:
            if not args.watch: raise
            print(f'Report refresh failed; preserving previous HTML: {type(exc).__name__}: {exc}',flush=True)
        if not args.watch or (series and all(l['phase']=='completed' for l in series['levels'])): break
        time.sleep(30)


if __name__ == '__main__': main()