#!/usr/bin/env python3
"""Generate brain-mesh.json and state-data.json for Three.js brain visualizer.

If nilearn is available, uses fsaverage5 mesh + Destrieux atlas.
Otherwise generates a simplified brain-shaped mesh using subdivided icosahedron.
"""

import json, math, os, sys

try:
    import numpy as np
    HAS_NUMPY = True
except ImportError:
    np = None
    HAS_NUMPY = False

try:
    from nilearn import datasets, surface
    HAS_NILEARN = True
except ImportError:
    HAS_NILEARN = False


# ─── Icosahedron subdivision ───────────────────────────────────────────────

G = (1 + math.sqrt(5)) / 2

ICOSA_VERTICES = [
    (-1, G, 0), (1, G, 0), (-1, -G, 0), (1, -G, 0),
    (0, -1, G), (0, 1, G), (0, -1, -G), (0, 1, -G),
    (G, 0, -1), (G, 0, 1), (-G, 0, -1), (-G, 0, 1),
]

ICOSA_FACES = [
    (0, 11, 5), (0, 5, 1), (0, 1, 7), (0, 7, 10), (0, 10, 11),
    (1, 5, 9), (5, 11, 4), (11, 10, 2), (10, 7, 6), (7, 1, 8),
    (3, 9, 4), (3, 4, 2), (3, 2, 6), (3, 6, 8), (3, 8, 9),
    (4, 9, 5), (2, 4, 11), (6, 2, 10), (8, 6, 7), (9, 8, 1),
]


def _midpoint(v1, v2, verts):
    x = (verts[v1][0] + verts[v2][0]) / 2
    y = (verts[v1][1] + verts[v2][1]) / 2
    z = (verts[v1][2] + verts[v2][2]) / 2
    length = math.sqrt(x*x + y*y + z*z)
    return (x/length, y/length, z/length)


def subdivide_icosahedron(detail):
    verts = list(ICOSA_VERTICES)
    faces = list(ICOSA_FACES)
    for _ in range(detail):
        mid_cache = {}
        new_faces = []
        for (a, b, c) in faces:
            ab = mid_cache.get((a, b)) or mid_cache.setdefault((a, b),
                len(verts) + len([k for k in mid_cache if k < (a, b)]))
            # We need a proper approach
        # Actually, let me use a simpler index tracking approach
    # Fall through to simpler implementation below


def generate_icosahedron_mesh(detail=4):
    """Generate subdivided icosahedron. detail=4 -> 2562 verts, detail=5 -> 10242."""
    verts = [(x/math.sqrt(x*x+y*y+z*z), y/math.sqrt(x*x+y*y+z*z), z/math.sqrt(x*x+y*y+z*z))
             for (x, y, z) in ICOSA_VERTICES]
    faces = list(ICOSA_FACES)

    for _ in range(detail):
        midpoint_map = {}
        new_faces = []
        for a, b, c in faces:
            key_ab = (a, b) if a < b else (b, a)
            key_bc = (b, c) if b < c else (c, b)
            key_ca = (c, a) if c < a else (a, c)
            if key_ab not in midpoint_map:
                midpoint_map[key_ab] = len(verts)
                verts.append(_midpoint(key_ab[0], key_ab[1], verts))
            if key_bc not in midpoint_map:
                midpoint_map[key_bc] = len(verts)
                verts.append(_midpoint(key_bc[0], key_bc[1], verts))
            if key_ca not in midpoint_map:
                midpoint_map[key_ca] = len(verts)
                verts.append(_midpoint(key_ca[0], key_ca[1], verts))
            ab = midpoint_map[key_ab]
            bc = midpoint_map[key_bc]
            ca = midpoint_map[key_ca]
            new_faces.append((a, ab, ca))
            new_faces.append((b, bc, ab))
            new_faces.append((c, ca, bc))
            new_faces.append((ab, bc, ca))
        faces = new_faces

    # Brain-like deformation
    deformed = []
    for (x, y, z) in verts:
        nx, ny, nz = x, y, z
        bump = 1 + 0.08 * math.sin(5*nx) * math.cos(4*ny) * math.sin(6*nz) \
                 + 0.05 * math.cos(8*nx*nx) * math.sin(7*ny*ny) \
                 + 0.03 * (0.5 + 0.5 * math.sin(10 * nz))
        front_stretch = 1 + 0.15 * (nx * nx if nx > 0 else 0)
        flat_bottom = 1 - 0.1 * max(0, -ny)
        scale = bump * front_stretch * flat_bottom
        deformed.append((x * scale, y * scale, z * scale))

    # Add fissure (split between hemispheres)
    # Push vertices near x=0, z>0 slightly apart
    fissured = []
    for (x, y, z) in deformed:
        if abs(x) < 0.15 and z > 0.1:
            x += 0.08 * (1 if z > 0 else -1)
        fissured.append((x, y, z))

    return {
        'vertices': fissured,
        'faces': faces,
    }


def load_nilearn_mesh():
    """Load fsaverage5 mesh using nilearn."""
    print("Loading fsaverage5 mesh via nilearn...")
    fsa = datasets.fetch_surf_fsaverage('fsaverage5')
    verts_l, faces_l = surface.load_surf_mesh(fsa['pial_left'])
    verts_r, faces_r = surface.load_surf_mesh(fsa['pial_right'])

    # Combine hemispheres
    verts = np.concatenate([verts_l, verts_r], axis=0)
    offset = verts_l.shape[0]
    faces = np.concatenate([faces_l, faces_r + offset], axis=0)

    return {
        'vertices': verts.tolist(),
        'faces': faces.tolist(),
    }


# ─── Region definitions ─────────────────────────────────────────────────────

# For simplified state data, define ROI functions on unit sphere
def _roi_weights(x, y, z):
    """Return dict of region weights for a vertex on unit sphere."""
    nx, ny, nz = x, y, z
    pfc = max(0, nx)**2
    parietal = max(0, ny)**2 * (1 - max(0, nx)**2)
    temporal = abs(nz)**1.5 * (1 - abs(nx))
    occipital = max(0, -nx)**2
    dmn = (0.3 * max(0, nx)**2 + 0.3 * max(0, -nx)**2 + 0.2 * abs(nz)**2)
    limbic = abs(nz)**2 * max(0, -ny + 0.2)**2
    return {'pfc': pfc, 'parietal': parietal, 'temporal': temporal,
            'occipital': occipital, 'dmn': dmn, 'limbic': limbic}


def generate_state_data_simple(verts_list):
    """Generate 4-state data using simplified ROI mapping."""
    # Use deterministic noise based on vertex index
    def _noise(i):
        return 0.5 + 0.5 * math.sin(i * 123.456 + 789.012)

    n = len(verts_list)
    states = {'baseline': [], 'focus12': [], 'focus15': [], 'trauma': []}

    for i, (x, y, z) in enumerate(verts_list):
        ln = math.sqrt(x*x + y*y + z*z)
        if ln == 0:
            w = {k: 0 for k in ('pfc','parietal','temporal','occipital','dmn','limbic')}
        else:
            w = _roi_weights(x/ln, y/ln, z/ln)

        nse = _noise(i)

        # Baseline: centered at 0.3 with noise, DMN slightly boosted
        v_base = max(0, min(1, 0.25 + 0.12 * nse + 0.1 * w['dmn']))

        # Focus 12: high PFC + parietal
        v_f12 = max(0, min(1, 0.1 + 0.7 * w['pfc'] + 0.5 * w['parietal'] + 0.1 * nse))

        # Focus 15: very high PFC + temporal + occipital
        v_f15 = max(0, min(1, 0.1 + 0.85 * w['pfc'] + 0.6 * w['temporal']
                           + 0.5 * w['occipital'] + 0.08 * nse))

        # Trauma: high limbic, low PFC, suppressed DMN
        v_tr = max(0, min(1, 0.15 + 0.7 * w['limbic']
                          + 0.05 * (1 - w['pfc'])
                          + 0.05 * w['dmn']
                          + 0.05 * nse))

        states['baseline'].append(v_base)
        states['focus12'].append(v_f12)
        states['focus15'].append(v_f15)
        states['trauma'].append(v_tr)

    return states


def generate_state_data_nilearn(verts_list):
    """Generate state data using Destrieux atlas regions via nilearn."""
    print("Generating state data with Destrieux atlas...")
    try:
        from nilearn import datasets, surface
        import numpy as np
    except ImportError:
        print("  nilearn import failed, falling back to simple mapping")
        return generate_state_data_simple(verts_list)

    verts = np.array(verts_list)
    n = verts.shape[0]

    try:
        destrieux = datasets.fetch_atlas_surf_destrieux()
        labels = destrieux['map_left']
        # Combine with right hemisphere
        labels = np.concatenate([labels, destrieux['map_right']])
        label_names = destrieux['labels']
    except Exception as e:
        print(f"  Destrieux atlas failed ({e}), using simple mapping")
        return generate_state_data_simple(verts_list)

    # Map Destrieux labels to functional groups
    # These are approximate mappings
    pfc_labels = set()
    parietal_labels = set()
    temporal_labels = set()
    occipital_labels = set()
    limbic_labels = set()
    dmn_labels = set()

    for idx, name in label_names:
        name_str = str(name).lower()
        if any(k in name_str for k in ['front', 'prefront', 'precent', 'fronto']):
            pfc_labels.add(idx)
        if any(k in name_str for k in ['pariet', 'supramarg', 'postcent']):
            parietal_labels.add(idx)
        if any(k in name_str for k in ['tempor', 'fusiform', 'parahippoc']):
            temporal_labels.add(idx)
        if any(k in name_str for k in ['occip', 'lingual', 'cuneus', 'calcar']):
            occipital_labels.add(idx)
        if any(k in name_str for k in ['amygdal', 'hippocamp', 'cingul', 'insul',
                                        'orbit', 'parolfact']):
            limbic_labels.add(idx)
        if any(k in name_str for k in ['cingul', 'precuneus', 'angular',
                                        'frontopol', 'temporopol']):
            dmn_labels.add(idx)

    # For vertices not in any known region, use fallback
    states = {'baseline': [], 'focus12': [], 'focus15': [], 'trauma': []}

    for i in range(n):
        lab = labels[i] if i < len(labels) else 0
        in_pfc = 1.0 if lab in pfc_labels else 0.0
        in_parietal = 1.0 if lab in parietal_labels else 0.0
        in_temporal = 1.0 if lab in temporal_labels else 0.0
        in_occipital = 1.0 if lab in occipital_labels else 0.0
        in_limbic = 1.0 if lab in limbic_labels else 0.0
        in_dmn = 1.0 if lab in dmn_labels else 0.0

        nse = 0.5 + 0.5 * math.sin(i * 123.456 + 789.012)

        # Add smooth falloff at region boundaries using spatial distance
        x, y, z = verts[i]
        ln = math.sqrt(x*x + y*y + z*z) or 1
        nx, ny, nz = x/ln, y/ln, z/ln

        spatial = _roi_weights(nx, ny, nz)

        # Blend label-based with spatial
        pfc_w = max(in_pfc, spatial['pfc'] * 0.3)
        par_w = max(in_parietal, spatial['parietal'] * 0.3)
        tmp_w = max(in_temporal, spatial['temporal'] * 0.3)
        occ_w = max(in_occipital, spatial['occipital'] * 0.3)
        lim_w = max(in_limbic, spatial['limbic'] * 0.3)
        dmn_w = max(in_dmn, spatial['dmn'] * 0.3)

        v_base = max(0, min(1, 0.25 + 0.12 * nse + 0.1 * dmn_w))
        v_f12 = max(0, min(1, 0.1 + 0.7 * pfc_w + 0.5 * par_w + 0.1 * nse))
        v_f15 = max(0, min(1, 0.1 + 0.85 * pfc_w + 0.6 * tmp_w + 0.5 * occ_w + 0.08 * nse))
        v_tr = max(0, min(1, 0.15 + 0.7 * lim_w + 0.05 * (1 - pfc_w) + 0.05 * dmn_w + 0.05 * nse))

        states['baseline'].append(v_base)
        states['focus12'].append(v_f12)
        states['focus15'].append(v_f15)
        states['trauma'].append(v_tr)

    return states


# ─── Main ───────────────────────────────────────────────────────────────────

def main():
    out_dir = os.path.dirname(os.path.abspath(__file__))

    # Generate mesh
    if HAS_NILEARN:
        try:
            mesh = load_nilearn_mesh()
            print(f"Loaded fsaverage5: {len(mesh['vertices'])} verts, {len(mesh['faces'])} faces")
        except Exception as e:
            print(f"nilearn mesh failed ({e}), using icosahedron fallback")
            mesh = generate_icosahedron_mesh(detail=4)
    else:
        print("nilearn not available. Generating simplified brain mesh.")
        mesh = generate_icosahedron_mesh(detail=4)

    mesh_path = os.path.join(out_dir, 'brain-mesh.json')
    with open(mesh_path, 'w') as f:
        json.dump({'vertices': mesh['vertices'], 'faces': mesh['faces']}, f)
    mesh_size = os.path.getsize(mesh_path)
    print(f"  brain-mesh.json: {mesh_size} bytes ({len(mesh['vertices'])} verts)")

    # Generate state data
    if HAS_NILEARN:
        state_data = generate_state_data_nilearn(mesh['vertices'])
    else:
        state_data = generate_state_data_simple(mesh['vertices'])

    # Round to 4 decimal places to reduce file size
    for state in state_data:
        state_data[state] = [round(v, 4) for v in state_data[state]]

    state_path = os.path.join(out_dir, 'state-data.json')
    with open(state_path, 'w') as f:
        json.dump(state_data, f)
    state_size = os.path.getsize(state_path)
    print(f"  state-data.json: {state_size} bytes ({len(state_data['baseline'])} vals per state)")

    print(f"\nDone. Output in: {out_dir}")
    print(json.dumps({
        'brain-mesh.json': mesh_size,
        'state-data.json': state_size,
    }))


if __name__ == '__main__':
    main()