/**
 * 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 { Mat3, Mat4, Vec3 } from '../mol-math/linear-algebra';
import { transformationForStateTransform } from './transformation';

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

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, pdbtmDescriptor: PDBTMDescriptor) {
    if (!MmcifFormat.is(model.sourceData)) return;

    let data = model.sourceData.data.db;
    let excludedChains: string[] = constructChainListFromOperations(pdbtmDescriptor);
    excludedChains = union(
        excludedChains,
        Array.from(data.pdbx_nonpoly_scheme.asym_id.toArray())
    );
    const structureOperations: StructureOperation[] = collectOperations(pdbtmDescriptor);
    console.log("Structure Operators:", structureOperations);

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

    const updated_struct_oper_list = createPdbxStructOperList(structureOperations, data.pdbx_struct_oper_list);
    console.log('Orig. struct oper list:', 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: updated_struct_oper_list
    });
}

function createPdbxStructAssemblyGen(
    structureOperations: StructureOperation[],
    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
    );
    const assembly_id: string[] = [];
    const asym_id_list: string[][] = [];
    const oper_expression: string[] = [];

    // we create only one assembly - maybe in multiple rows
    assembly_id.push("1");
    asym_id_list.push(asym_id_list_column);
    oper_expression.push("0"); // this is the basic membrane plane transformation

    const chainMap = new Map<string, string[]>();
    // collect each operator for each chain
    structureOperations.slice(1).forEach(function (operation: StructureOperation) {
        if (!chainMap.has(operation.chainId)) {
            chainMap.set(operation.chainId, [ operation.id.toString() ]);
        } else {
            chainMap.get(operation.chainId)?.push(operation.id.toString());
        }
    });
    // fill columns of table
    chainMap.forEach(function(operators, chain) {
        assembly_id.push("1");
        asym_id_list.push([ chain ]);
        oper_expression.push(operators.join(","));
    });


    // create table with new column
    let updated_pdbx_struct_assembly_gen = Table.ofColumns(
        pdbx_struct_assembly_gen._schema,
        {
            assembly_id: Column.ofStringArray(assembly_id),
            asym_id_list: Column.ofStringListArray(asym_id_list),
            oper_expression: Column.ofStringArray(oper_expression)
        }
    );
    console.log('Orig. assembly_gen', Table.formatToString(pdbx_struct_assembly_gen));
    console.log('Updated assembly_gen', Table.formatToString(updated_pdbx_struct_assembly_gen));

    return updated_pdbx_struct_assembly_gen;
}

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

    const id: string[] = [];
    const type: string[] = [];
    const name: string[] = [];
    const symmetry_operation: string[] = [];
    const matrix: number[][] = [];
    const vector: number[][] = [];

    structureOperations.forEach(function (operation: StructureOperation) {
        id.push(operation.id.toString());
        type.push(".");
        const operationName = (operation.id == 0)
            ? "Membrane transformation as identity operation"
            : `${operation.chainId} -> ${operation.newChainId}`;
        name.push(operationName);
        symmetry_operation.push(".");

        const rotation: number[] = Mat3.toArray(operation.operator.rotation, [], 0);
        const translation: number[] = Vec3.toArray(operation.operator.translation, [], 0);;
        matrix.push(rotation);
        vector.push(translation);
    });


    let updated_pdbx_struct_oper_list = Table.ofColumns(
        pdbx_struct_oper_list._schema,
        {
            id: Column.ofStringArray(id),
            type: Column.ofArray({
                array: type,
                schema: pdbx_struct_oper_list.type.schema
            }),
            name: Column.ofStringArray(name),
            symmetry_operation: Column.ofStringArray(symmetry_operation),
            matrix:  Column.ofArray({
                array: matrix,
                schema: pdbx_struct_oper_list.matrix.schema
            }),
            vector: Column.ofArray({
                array: vector,
                schema: pdbx_struct_oper_list.vector.schema
            })
        }
    );

    console.log('Orig. pdbx_struct_oper_list', Table.formatToString(pdbx_struct_oper_list));
    console.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);
    console.log('Excluded chains:', excludedChains);
    console.log('Included chains:', asym_id_list);

    return 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());
}


//////////////////////////////////////////////////////////////////////


/*
 * Construction of struct operators
 */

/**
 * Type to contain operation properties.
 */
type StructureOperation = {
    id: number,
    chainId: string,
    newChainId: string,
    operator: {
        rotation: Mat3,
        translation: Vec3
    }
}

function collectOperations(descriptor: PDBTMDescriptor): StructureOperation[] {

    const operations: StructureOperation[] = [];
    const annotations = descriptor.additional_entry_annotations;

    const membraneTransformation = transformationForStateTransform(descriptor.additional_entry_annotations.membrane.transformation_matrix);
    operations.push({
        id: 0,
        chainId: "Any chain with identity operation",
        newChainId: ".",
        operator: mat4ToRotationAndTranslation(
            Mat4.mul(Mat4(), Mat4.rotX90, membraneTransformation)
        )
    });

    let operationId: number = 1;

    if (annotations?.biomatrix?.matrix_list) {
        annotations.biomatrix.matrix_list.forEach(function(mx) {
            mx.apply_to_chain_list.forEach(function(chainPair) {
                let id  = chainPair.chain_id;
                let newId = chainPair.new_chain_id;
                if (annotations.biomatrix.chain_deletes?.includes(newId)) {
                    console.log(`${id} -> ${newId} transformation skipped due to delete rule`);
                    return;
                }
                const mtx = transformationForStateTransform(mx.transformation_matrix);
                const composedTransformation = Mat4.mul(
                    Mat4(),
                    Mat4.rotX90,
                    Mat4.mul(Mat4(), membraneTransformation, mtx)
                );

                const operator: StructureOperation = {
                    id: operationId,
                    chainId: id,
                    newChainId: newId,
                    operator: mat4ToRotationAndTranslation(composedTransformation)
                };
                operations.push(operator);
                operationId++;
            });
        });
    }

    return operations;
}

/**
 * Convert Mat4 object to array representation of its 3x4 part.
 * Rotation matrix 3x3 and translation vector 3x1.
 *
 * @param mat4 Mat4 object
 * @returns Object with rotation matrix and translation vector
 */
function mat4ToRotationAndTranslation(mat4: Mat4) {
    return {
        rotation: Mat3.fromMat4(
            Mat3.zero(), Mat4.extractRotation(Mat4.zero(), mat4)
        ),
        translation: Mat4.getTranslation(Vec3.zero(), mat4)
    };
}