TmMolStar/src/tmdet-extension/tmdet-color-theme.ts

/**
 * Copyright (C) 2022, Protein Bioinformatics Research Group, 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 { StructureElement, Unit } from 'molstar/lib/mol-model/structure';
import { ColorTheme } from 'molstar/lib/mol-theme/color';
import { ThemeDataContext } from 'molstar/lib/mol-theme/theme';
import { Color } from 'molstar/lib/mol-util/color';
import { ColorNames } from 'molstar/lib/mol-util/color/names';
import { ParamDefinition as PD } from 'molstar/lib/mol-util/param-definition';
import { Location } from 'molstar/lib/mol-model/location';
import { ChainList, getResidue, PDBTMChain } from './types';
import { TmDetChainListCache, TmDetDescriptorCache } from './prop';


export type TmDetColorThemeParams = {
    pdbtmDescriptor: any
};

export function TmDetColorTheme(
    ctx: ThemeDataContext,
    props: PD.Values<TmDetColorThemeParams>
): ColorTheme<TmDetColorThemeParams> {
    let descriptorChains: PDBTMChain[] = [];
    let pdbtmDescriptor = props.pdbtmDescriptor;
    const pdbId = ctx.structure?.model.entryId;

    // If it is not given as parameter,
    // but there is one in the cache for this structure.
    if (!pdbtmDescriptor) {
        if (pdbId) {
            pdbtmDescriptor = TmDetDescriptorCache.get(pdbId);
        }
    }
    if (pdbtmDescriptor) {
        descriptorChains = pdbtmDescriptor.chains;
    }

    const tmType = pdbtmDescriptor.additional_entry_annotations.tm_type;
    const chainList = createResidueListsPerChain(descriptorChains, pdbtmDescriptor.side1, tmType);
    if (pdbId) {
        TmDetChainListCache.set(pdbId, chainList);
    }

    return {
        factory: TmDetColorTheme,
        granularity: 'group', //'residue' as any,
        color: (location: Location) => getColor(location, chainList),
        props: props,
        description: 'TMDet...',
    };
}

const DefaultResidueColor = ColorNames.lightgrey;

enum SiteIndexes {
    Side1 = 0,
    Side2 = 1,
    TmAlpha = 2,
    TmBeta = 3,
    TmReentrantLoop = 4,
    InterfacialHelix = 5,
    UnknownLocalization = 6,
    MembraneInside = 7,
    Periplasm = 8
};

// Old default values - it is overwritten by ult_* CSS classes
// See below updateSiteColors().
const siteColors = [
    Color.fromArray([255,100,100], 0), // Side1
    Color.fromArray([100,100,255], 0), // Side2
    Color.fromArray([255,255,  0], 0), // TM alpha
    Color.fromArray([255,255,  0], 0), // TM beta
    Color.fromArray([255,127,  0], 0), // TM re-entrant loop
    Color.fromArray([0,255,    0], 0), // Interfacial Helix
    Color.fromArray([196,196,196], 0), // Unknow localization
    Color.fromArray([0,255,    0], 0), // Membrane Inside
    Color.fromArray([255, 0, 255], 0)  // Periplasm
];

const siteCssNames = [
    "ult_side1",
    "ult_side2",
    "ult_alpha",
    "ult_beta",
    "ult_reentrant",
    "ult_ifh",
    "ult_unknown",
    "ult_membins",
    "ult_periplasm"
];

const regionColorMapFromCss = new Map();

// set default values
siteCssNames.forEach((className, index) => {
    regionColorMapFromCss.set(className, siteColors[index]);
});

function getColor(location: Location, chains: ChainList): Color {
    let color = DefaultResidueColor;

    // TODO: How to solve cases when chain operation occurs?
    if (StructureElement.Location.is(location) && Unit.isAtomic(location.unit)) {

        const { chainId, residueId } = getChainAndResidueIds(location);
        color = residueColor(chains, chainId!, residueId!);
    }
    return color;
}

export function getChainAndResidueIds(location: Location) {
    let chainId = undefined;
    let residueId = undefined;

    if (StructureElement.Location.is(location) && Unit.isAtomic(location.unit)) {
        const atomicHierarchy = location.unit.model.atomicHierarchy;
        const residueIdx = StructureElement.residueIndex(location);
        const chainIdx = StructureElement.Location.chainIndex(location);

        residueId = atomicHierarchy.residues.label_seq_id.value(residueIdx).toString();
        chainId = atomicHierarchy.chains.label_asym_id.value(chainIdx).toString();
    }
    return {
        chainId: chainId,
        residueId: residueId
    };
}

export function createResidueListsPerChain(chains: PDBTMChain[], side1: string|null, tmType: string) {

    const chainList: ChainList = [];
    const hasBeta = chains.some(chain => chain.additional_chain_annotations.type == "beta");

    chains.forEach((chain: PDBTMChain) => {
        const chainType = chain.additional_chain_annotations.type;
        chainList.push({ chainId: chain.chain_label, type: chainType, residues: [] });
        let annotationErrorLogged = false;
        chain.residues.forEach((residue) => {
            let siteId = residue.site_data![0].site_id_ref - 1;
            let siteColorId = siteId;
            if (tmType == "Tm_Alpha" && hasBeta && side1 == "Periplasm") {
                if (siteId == SiteIndexes.Side1) {
                    siteColorId = SiteIndexes.Periplasm;
                } else if (siteId == SiteIndexes.Side2) {
                    if (chainType == "non_tm" || chainType == "alpha") {
                        siteColorId = SiteIndexes.Side1;
                    } else if (chainType == "beta") {
                        siteColorId = SiteIndexes.Side2;
                    }
                }
            } else if (chainType == "beta") {
                if (side1 == "Periplasm" && siteColorId == SiteIndexes.Side1
                    || side1 == "Outside" && siteColorId == SiteIndexes.Side2) {
                    siteColorId = SiteIndexes.Periplasm;
                } else if (!annotationErrorLogged && side1 == "Inside" && siteColorId == SiteIndexes.Side1) {
                    console.error(`Annotation error: beta chain has inside region ${chain.chain_label}:${residue.pdb_res_label}`);
                    annotationErrorLogged = true;
                }
            }

            chainList[chainList.length - 1].residues.push({
                authId: residue.pdb_res_label,
                siteId: siteId,
                siteColorId: siteColorId
            });
        });
    });

    return chainList;
}

function residueColor(chains: ChainList, chainId: string, residueId: string): Color {
    let color = DefaultResidueColor;

    const residue = getResidue(chains, chainId, residueId);
    if (residue) {
        color = siteColors[residue.siteColorId];
    }

    return color;
}

// Provider

export const TmDetColorThemeProvider: ColorTheme.Provider<TmDetColorThemeParams, 'tmdet-custom-color-theme'> = {
    name: 'tmdet-custom-color-theme',
    label: 'TMDet Topology Theme',
    category: 'TMDet',
    factory: TmDetColorTheme,
    getParams: () => ({ pdbtmDescriptor: { isOptional: true } }),
    defaultValues: { pdbtmDescriptor: undefined },
    isApplicable: () => true,
};



// Colors from CSS rules

function loadRegionColorsFromStyleSheets(prefix: string = 'ult_'): void {
    const sheets: CSSStyleSheet[] = Array.from(document.styleSheets);
    sheets.forEach((sheet: CSSStyleSheet) => {
        const rules: CSSRule[] = Array.from(sheet.cssRules);
        rules.forEach((rule: CSSRule) => fetchRule(rule, prefix));
    });
}

export function updateSiteColors(side1: "Inside"|"Outside"|null): void {
    if (!side1) {
        return;
    }
    loadRegionColorsFromStyleSheets();
    if (regionColorMapFromCss.size == 0) {
        console.warn('Cannot read any region-related color rules');
    }

    siteCssNames.forEach((ultClassName, index) => {
        const color = regionColorMapFromCss.get(ultClassName);
        if (color != null) {
            siteColors[index] = color;
        }
    });

    const inside = regionColorMapFromCss.get("ult_inside");
    const outside = regionColorMapFromCss.get("ult_outside");
    if (side1 == "Inside") {
        siteColors[SiteIndexes.Side1] = inside;
        siteColors[SiteIndexes.Side2] = outside;
    } else if (side1 == "Outside") {
        siteColors[SiteIndexes.Side1] = outside;
        siteColors[SiteIndexes.Side2] = inside;
    }
}

function fetchRule(rule: CSSRule, prefix: string) {
    let styleRule = rule as CSSStyleRule;
    if (styleRule.selectorText?.startsWith('.' + prefix)) {
        const value = styleRule.style.getPropertyValue('fill');
        const color: Color = getStyleColor(value);
        const key = styleRule.selectorText.slice(1);
        regionColorMapFromCss.set(key, color);
    }
}

function getStyleColor(cssColorText: string): Color {
    const values = cssColorText?.match(/\d+/g);
    let intValues = values?.map(value => parseInt(value));
    if (!intValues) {
        intValues = [ 0, 0, 0 ];
    }
    return Color.fromArray(intValues, 0);
}