/**
 * 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 { MmcifFormat } from '../../mol-model-formats/structure/mmcif';
import { Column, Table } from '../../mol-data/db';
import { mmCIF_Schema } from '../../mol-io/reader/cif/schema/mmcif';
import { Model } from '../../mol-model/structure';
import { ModelSymmetry } from '../../mol-model-formats/structure/property/symmetry';
import { PDBTMDescriptor } from './types';
import { DebugUtil } from './debug-utils';


export function registerTmDetSymmetry(pdbtmDescriptor: PDBTMDescriptor) {
    ModelSymmetry.Provider.formatRegistry.remove('mmCIF');
    const excludedChains = constructChainListFromOperations(pdbtmDescriptor);
    ModelSymmetry.Provider.formatRegistry.add('mmCIF',  function(model: Model) {
        return tmDetSymmetryFromMmCif(model, excludedChains);
    });
}

function constructChainListFromOperations(pdbtmDescriptor: PDBTMDescriptor): string[] {
    const excludedChains: string[] = [];
    // add chain deletes
    const biomatrix = pdbtmDescriptor.additional_entry_annotations.biomatrix;
    if (biomatrix?.chain_deletes) {
        biomatrix.chain_deletes.forEach(
            chainId => excludedChains.push(chainId)
        );
    }
    // exclude result of transformations
    if (biomatrix?.matrix_list) {
        biomatrix.matrix_list.forEach(
            matrix => matrix.apply_to_chain_list.forEach(
                applyItem => excludedChains.push(applyItem.new_chain_id)
            )
        );
    }

    return excludedChains;
}

function tmDetSymmetryFromMmCif(model: Model, excludedChains: string[]) {
    if (!MmcifFormat.is(model.sourceData)) return;

    let data = model.sourceData.data.db;
    excludedChains = union(
        excludedChains,
        Array.from(data.pdbx_nonpoly_scheme.asym_id.toArray())
    );

    const updated_pdbx_struct_assembly_gen = createPdbxStructAssemblyGen(
        data.pdbx_struct_assembly_gen,
        excludedChains
    );
    DebugUtil.log('Non-poly entities:', Table.formatToString(data.pdbx_entity_nonpoly));
    DebugUtil.log('Non-poly chains:', data.pdbx_nonpoly_scheme.asym_id.toArray());

    const only_identity_operation = createPdbxStructOperList(data.pdbx_struct_oper_list);

    return ModelSymmetry.fromData({
        symmetry: data.symmetry,
        cell: data.cell,
        struct_ncs_oper: data.struct_ncs_oper,
        atom_sites: data.atom_sites,
        pdbx_struct_assembly: data.pdbx_struct_assembly,
        pdbx_struct_assembly_gen: updated_pdbx_struct_assembly_gen,
        pdbx_struct_oper_list: only_identity_operation
    });
}

function createPdbxStructAssemblyGen(pdbx_struct_assembly_gen: Table<mmCIF_Schema['pdbx_struct_assembly_gen']>,
    excludedChains: string[]): Table<mmCIF_Schema['pdbx_struct_assembly_gen']> {

    const asym_id_list_column = createAsymIdColumn(
        pdbx_struct_assembly_gen, excludedChains
    );

    // create table with new column
    let updated_pdbx_struct_assembly_gen = Table.ofColumns(
        pdbx_struct_assembly_gen._schema,
        {
            assembly_id: Column.ofStringArray([ '1' ]),
            asym_id_list: asym_id_list_column,
            //oper_expression: data.pdbx_struct_assembly_gen.oper_expression
            // NOTE: we expect here pdbx_struct_assembly_gen has only one row
            oper_expression: Column.ofStringArray([ '1' ])
        }
    );
    DebugUtil.log('Orig. assembly_gen', Table.formatToString(pdbx_struct_assembly_gen));
    DebugUtil.log('Updated assembly_gen', Table.formatToString(updated_pdbx_struct_assembly_gen));

    return updated_pdbx_struct_assembly_gen;
}

function createPdbxStructOperList(pdbx_struct_oper_list: Table<mmCIF_Schema['pdbx_struct_oper_list']>):
    Table<mmCIF_Schema['pdbx_struct_oper_list']> {

    let updated_pdbx_struct_oper_list = Table.ofColumns(
        pdbx_struct_oper_list._schema,
        {
            id: Column.ofStringArray([ '1' ]),
            type: Column.ofArray({
                array: [ 'identity operation' ],
                schema: pdbx_struct_oper_list.type.schema
            }),
            name: Column.ofStringArray([ '1_555' ]),
            symmetry_operation: Column.ofStringArray([ 'x,y,z' ]),
            matrix:  Column.ofArray({
                array: [[ 1, 0, 0, 0, 1, 0, 0, 0, 1 ]],
                schema: pdbx_struct_oper_list.matrix.schema
            }),
            vector: Column.ofArray({
                array: [ [ 0, 0, 0 ] ],
                schema: pdbx_struct_oper_list.vector.schema
            })
        }
    );

    DebugUtil.log('Orig. pdbx_struct_oper_list', Table.formatToString(pdbx_struct_oper_list));
    DebugUtil.log('Updated pdbx_struct_oper_list', Table.formatToString(updated_pdbx_struct_oper_list));

    return updated_pdbx_struct_oper_list;
}

function createAsymIdColumn(pdbx_struct_assembly_gen: Table<mmCIF_Schema['pdbx_struct_assembly_gen']>,
    excludedChains: string[]) {

    let asym_id_list: string[] = [];
    for (let i = 0; i < pdbx_struct_assembly_gen._rowCount; i++) {
        const currentAsymIdList = pdbx_struct_assembly_gen.asym_id_list.value(i);
        asym_id_list = asym_id_list.concat(currentAsymIdList);
    }
    asym_id_list = minus(asym_id_list, excludedChains);
    DebugUtil.log('Excluded chains:', excludedChains);
    DebugUtil.log('Included chains:', asym_id_list);

    return Column.ofStringListArray([ asym_id_list ]);
}

// difference of two string arrays (interpreted as sets)
function minus(a: string[], b: string[]): string[] {
    const b_set = new Set(b);
    const difference = a.filter(x => !b_set.has(x));
    return Array.from(new Set(difference).values());
}

// union of two string arrays (interpreted as sets)
function union(a: string[], b: string[]): string[] {
    const a_set = new Set(a);
    b.forEach(item => a_set.add(item));
    return Array.from(a_set.values());
}