TmMolStar/src/Tm3DStandaloneApp/tmdet-viewer/TmViewerStandalone.ts

/**
 * Copyright (c) 2019-2020 mol* contributors, licensed under MIT, See LICENSE file for more info.
 *
 * @author Alexander Rose <alexander.rose@weirdbyte.de>
 * @author Joan Segura <joan.segura@rcsb.org>
 * @author Yana Rose <yana.rose@rcsb.org>
 * @author Sebastian Bittrich <sebastian.bittrich@rcsb.org>
 */
/**
 * Copyright (C) 2024, Protein Bioinformatics Research Group, HUN-REN RCNS
 *
 * Licensed under CC BY-NC 4.0, see LICENSE file for more info.
 *
 * @author Gabor Tusnady <tusnady.gabor@ttk.hu>
 * @author Csongor Gerdan <gerdan.csongor@ttk.hu>
 */

import { Viewer, ViewerProps } from '@rcsb/rcsb-molstar/build/src/viewer';
import { TmDetColorThemeProvider, updateSiteColors } from '../../TmFv3DApp/tmdet-extension/tmdet-color-theme';
import { MembraneOrientationPreset } from '../../TmFv3DApp/tmdet-extension/behavior';
import { TmDetLabelProvider } from '../../TmFv3DApp/tmdet-extension/labeling';
import { BuiltInTrajectoryFormat } from 'molstar/lib/mol-plugin-state/formats/trajectory';
import { PresetProps } from '@rcsb/rcsb-molstar/build/src/viewer/helpers/preset';
import { Mat4 } from 'molstar/lib/commonjs/mol-math/linear-algebra';
import { TrajectoryHierarchyPresetProvider } from 'molstar/lib/mol-plugin-state/builder/structure/hierarchy-preset';
import { registerRegionDescriptorData } from '../../TmFv3DApp/UniTmpHelper';
import { TmDetRcsbPreset } from '../TmTrajectoryHierarchyPreset';
import { MolScriptBuilder } from 'molstar/lib/mol-script/language/builder';
import { StateTransforms } from 'molstar/lib/mol-plugin-state/transforms';
import { Color } from 'molstar/lib/mol-util/color';

import { Script } from 'molstar/lib/mol-script/script';
import { ColorNames } from 'molstar/lib/mol-util/color/names';

export class TmViewerStandalone extends Viewer {

    constructor(elementOrId: string | HTMLElement, props: Partial<ViewerProps> = {}, withTmDetBehaviour: boolean = true) {

        super(elementOrId, props);

        if (withTmDetBehaviour) {
            this.initBehaviour();
        }
    }

    public async initBehaviour() {
        let tree = this.plugin.state.behaviors.build();


        this.plugin.builders.structure.representation.registerPreset(MembraneOrientationPreset);
        this.plugin.representation.structure.themes.colorThemeRegistry.add(TmDetColorThemeProvider);
        this.plugin.managers.lociLabels.addProvider(TmDetLabelProvider);

        await this.plugin.runTask(this.plugin.state.behaviors.updateTree(tree, { doNotUpdateCurrent: true, doNotLogTiming: true }));
    }

    loadStructureFromUrl<P, S>(url: string, format: BuiltInTrajectoryFormat, isBinary: boolean, config?: {
        props?: PresetProps & {
            dataLabel?: string;
        };
        matrix?: Mat4;
        reprProvider?: TrajectoryHierarchyPresetProvider<P, S>;
        params?: P;
    }): Promise<any> {
        return super.loadStructureFromUrl(url, format, isBinary, config);
    }

    chainSelection(auth_asym_id: string) {
        return MolScriptBuilder.struct.generator.atomGroups({
            'chain-test': MolScriptBuilder.core.rel.eq([MolScriptBuilder.struct.atomProperty.macromolecular.auth_asym_id(), auth_asym_id])
        });
    }

    setChainColor(labelAsymIds: string[], hexRgb: string|false = false) {
        this.overPaint(labelAsymIds, hexRgb, false);
    }

    resetChainColor(labelAsymIds: string[]) {
        this.overPaint(labelAsymIds, false, true);
    }

    /**
     * Util function to handle overpaint logic.
     * @param labelAsymId 
     * @param hexRgb default color grey, when it is false
     * @param clear if true, it clears the overpaint color
     */
    protected overPaint(labelAsymIds: string[], hexRgb: string|false = false, clear: boolean = false) {
        if (labelAsymIds.length == 0) {
            console.warn('No specified chain(s) for overpaintings');
            return;
        }

        const state = this.plugin.state.data;

        const selected = state.selectQ(q => q.ofTransformer(StateTransforms.Representation.StructureRepresentation3D));
        if (selected.length == 0) {
            console.log('Representation not found');
            return;
        }

        const chainList = labelAsymIds.join(' ');
        selected.forEach(async (representation) => {
            const update = this.plugin.build().to(representation);
            update.apply(StateTransforms.Representation.OverpaintStructureRepresentation3DFromScript, {
                layers: [
                    {
                        script:  Script(
                            `(sel.atom.atom-groups :chain-test (set.has (set ${chainList}) atom.label_asym_id))`,
                            'mol-script'
                        ),
                        color:  hexRgb !== false ? Color.fromHexString(hexRgb) : ColorNames.white,
                        clear: clear
                    }
                ]
            });
            await update.commit();
        });
    }

    async loadWithUNITMPMembraneRepresentation(params: any) {
        setTimeout(() => { this.plugin.clear(); }, 100); // clear scene after some delay
        updateSiteColors(params.side1);

        setTimeout(() => { (async () => {
            await registerRegionDescriptorData(params.regionDescriptorUrl, params.side1);

            // load structure
            this.loadStructureFromUrl(params.structureUrl, 'mmcif', false, { reprProvider: TmDetRcsbPreset });

        })(); }, 500);
    }

}