// Adapted from aaredb (gitea.psi.ch/mx/aaredb, the PSI MX group's own code), frontend/src/components/Structure/molstar/loadMaps.ts (+ the anomalous layer) at 8ece2c2.
import { PluginContext } from 'molstar/lib/mol-plugin/context';
import { StateTransforms } from 'molstar/lib/mol-plugin-state/transforms';
import { createVolumeRepresentationParams } from 'molstar/lib/mol-plugin-state/helpers/volume-representation-params';
import { Volume } from 'molstar/lib/mol-model/volume';
import { Color } from 'molstar/lib/mol-util/color';
import { setSubtreeVisibility } from 'molstar/lib/mol-plugin/behavior/static/state';
import type { ComputedMap } from './densityCrop';

// Density volumes in the Mol* state tree. σ sliders map user-σ to ABSOLUTE iso
// values via the map's full-grid mean/rms (reported by the worker) — Mol*'s own
// relative sigma would re-derive stats from whatever grid it holds, which stops
// being crystallographically meaningful the moment we crop. A σ change is a
// representation param update (re-mesh only — no FFT, no re-upload).

export const MAP_COLORS = {
    twofofc: 0x4d94d6, // steel-blue
    fofcPos: 0x35b96a, // green
    fofcNeg: 0xd6564d, // red
    anom: 0xe040e0, // magenta (anomalous difference map, +sigma only)
} as const;

export interface MapLayer {
    volumeRef: string;
    reprRefs: string[]; // one for 2fofc and anom, two (+/−) for fofc
    rmsd: number;
    mean: number;
}

const abs = (map: { mean: number; rmsd: number }, sigma: number) =>
    Volume.IsoValue.absolute(map.mean + sigma * map.rmsd);

async function addIso(
    plugin: PluginContext,
    volumeRef: string,
    map: ComputedMap,
    sigma: number,
    color: number,
): Promise<string> {
    const volumeCell = plugin.state.data.select(volumeRef)[0];
    const repr = plugin
        .build()
        .to(volumeRef)
        .apply(
            StateTransforms.Representation.VolumeRepresentation3D,
            createVolumeRepresentationParams(plugin, volumeCell?.obj?.data, {
                type: 'isosurface',
                typeParams: { isoValue: abs(map, sigma), visuals: ['wireframe'], sizeFactor: 1 },
                color: 'uniform',
                colorParams: { value: Color(color) },
            }),
        );
    await repr.commit();
    // The mesh must NOT be pickable: a pickable isosurface shields the model,
    // so residue clicks and middle-click recenter would hit map cells instead
    // of atoms. Pickability is runtime state on the INNER representation
    // (cell.obj.data.repr — the cell data itself has no setState), and the
    // pick buffer needs a redraw to drop the mesh.
    const reprCell = plugin.state.data.select(repr.ref)[0];
    const reprObj = (reprCell?.obj?.data as { repr?: { setState?: (s: { pickable: boolean }) => void } })?.repr;
    reprObj?.setState?.({ pickable: false });
    plugin.canvas3d?.requestDraw?.();
    return repr.ref;
}

export async function loadMapLayer(
    plugin: PluginContext,
    map: ComputedMap,
    kind: '2fofc' | 'fofc' | 'anom',
    sigma: number,
): Promise<MapLayer> {
    const data = await plugin.builders.data.rawData({
        data: new Uint8Array(map.ccp4),
        label: kind,
    });
    const parsed = await plugin.dataFormats.get('ccp4')!.parse(plugin, data);
    const volumeRef: string = parsed.volume.ref;

    const reprRefs: string[] = [];
    if (kind === '2fofc') {
        reprRefs.push(await addIso(plugin, volumeRef, map, sigma, MAP_COLORS.twofofc));
    } else if (kind === 'anom') {
        reprRefs.push(await addIso(plugin, volumeRef, map, sigma, MAP_COLORS.anom));
    } else {
        reprRefs.push(await addIso(plugin, volumeRef, map, +sigma, MAP_COLORS.fofcPos));
        reprRefs.push(await addIso(plugin, volumeRef, map, -sigma, MAP_COLORS.fofcNeg));
    }
    return { volumeRef, reprRefs, rmsd: map.rmsd, mean: map.mean };
}

/** σ slider update: param-only change, re-mesh without recompute. */
export async function updateLayerSigma(
    plugin: PluginContext,
    layer: MapLayer,
    sigma: number,
): Promise<void> {
    const sigmas = layer.reprRefs.length === 2 ? [+sigma, -sigma] : [sigma];
    const b = plugin.build();
    layer.reprRefs.forEach((ref, i) => {
        b.to(ref).update((old: { type: { params: { isoValue: unknown } } }) => {
            old.type.params.isoValue = abs(layer, sigmas[i]);
        });
    });
    await b.commit();
}

export function setLayerVisible(plugin: PluginContext, layer: MapLayer, visible: boolean): void {
    setSubtreeVisibility(plugin.state.data, layer.volumeRef, !visible); // arg = isHidden
}