import _, { chain, differenceBy, keyBy, max, zip, zipWith } from 'lodash';
import Selection from './models/selection';
import Protein from "./models/protein";
import { Segmentation, SegmentationType } from "./models/segmentation";
import Loop from "./models/loop";
import wait_for_computation_async from "../async_calls"


const toBase64 = file => new Promise((resolve, reject) => {
    const reader = new FileReader();
    reader.readAsDataURL(file);
    reader.onload = () => resolve(reader.result);
    reader.onerror = error => reject(error);
});

function roundToTwo(num) {
    return +(Math.round(num + "e+2") + "e-2");
}


export default {
    state: {
        proteins: {},
        selections: {},
        bfactors: {},
        currentSelection: new Selection(),
        threeDProtein: null,
        threeDFlexibility: null,
        zoomedLoop: null,
        tooltipText: "",
        tooltipShown: false,

    },
    getters: {
        getProteins(state, getters, rootState) {
            return state.proteins;
        },
        getProtein: (state) => (proteinName) => state.proteins[proteinName],

        getThreeDProtein(state, getters, rootState) {
            return state.threeDProtein;
        },


        getSelections(state, getters, rootState) {
            return state.selections;
        },

        getCurrentSelection(state, getters, rootState) {
            return state.currentSelection;
        },

        getBFactors(state, getters, rootState) {
            return state.bfactors;
        },

        getBFactorsByProtein: (state, getters, rootState) => (proteinName) => {
            return state.bfactors[proteinName];
        },

        getAllLoops: (state, getters) => {
            let loops = []
            loops.push(...getters.getScaffold.loops, ...getters.getInsert.loops);

            return loops;
        },

        getZoomedLoop(state) {
            return state.zoomedLoop;
        },
        getTooltipText(state) {
            return state.tooltipText;
        },
        getTooltipShown(state) {
            return state.tooltipShown;
        },
        get3dView(state) {
            return {
                loop: state.zoomedLoop,
                protein: state.threeDProtein,
                flexibility: state.threeDFlexibility
            }
        },

        getMaxCcorWithActive: (state, getters) => (loop_id, strategy) => {
            let loop = getters.getScaffold.loops.find(l => l.id == loop_id);
            let loops_of_interest = getters.getWantedLoops;
            let ccor_values = []
            for (let interesting_loop_id of loops_of_interest.map(l => l.id)) {
                if (interesting_loop_id != loop_id) {
                    ccor_values.push(Math.abs(strategy(loop, interesting_loop_id)))
                }
            }
            return max(ccor_values)
        },

        getScaffold(state) {
            return state.proteins[state.scaffold];
        },

        getInsert(state) {
            return state.proteins[state.insert];
        },

        getRole: (state, getters) => (role) => {
            return role == "scaffold" ? getters.getScaffold : getters.getInsert
        }


    },
    mutations: {
        addProtein(state, protein) {
            let newProteins = { ...state.proteins };
            newProteins[protein.name] = protein;
            state.proteins = newProteins;
        },


        setScaffold(state, proteinName) {
            state.scaffold = proteinName;
        },

        setInsert(state, proteinName) {
            state.insert = proteinName;
        },

        setThreeDProtein(state, { protein, flexibility }) {
            state.threeDProtein = protein;
            state.threeDFlexibility = flexibility;
            state.zoomedLoop = null;
        },

        removeProtein(state, name) {
            let newProteins = { ...state.proteins };
            state.proteins = _.omit(newProteins, name);
        },

        addBFactors(state, { proteinName, bfactors }) {
            let newBfactors = { ...state.bfactors };
            newBfactors[proteinName] = bfactors;
            state.bfactors = newBfactors;
        },

        removeBFactors(state, { proteinName }) {
            let newBfactors = { ...state.bfactors };
            state.bfactors = _.omit(newBfactors, proteinName);
        },

        addCustomLoop(state, { proteinName, loop }) {
            state.proteins[proteinName].customLoops.push(loop);
        },
        addComuptedLoop(state, { proteinName, loop }) {
            state.proteins[proteinName].computedLoops.push(loop);
        },
        addComputedSuperLoop(state, { proteinName, loop }) {
            state.proteins[proteinName].computedSuperloops.push(loop);
        },


        completeSelection(state) {
            // let newSelections = { ...state.selections };
            // newSelections[state.currentSelection.protein.name + "_" + state.currentSelection.from.pdb_number + "_" + state.currentSelection.to.pdb_number] = state.currentSelection;
            // state.selections = newSelections;
            state.currentSelection = new Selection();
        },

        select(state, { protein, acid }) {
            state.currentSelection.select(protein, acid);
        },

        cancelSelection(state) {
            state.currentSelection = new Selection();
        },

        zoomToLoop(state, loop) {
            state.zoomedLoop = loop;
        },
        setLoopLocation(state, { proteinName, loop_id, location }) {
            let protein = state.proteins[proteinName];
            let loop = protein.loops.find(l => l.id == loop_id);

            let before = { start: loop.start, end: loop.end };

            if (before.start < location.start) {
                for (let acid of protein.getAcids(before.start, location.start - 1)) {
                    acid.secondary_structure = 'C';
                }
                protein.invalidateLoopsGeometry(before.start, location.start)
            } else if (location.start < before.start) {
                let newSecs = protein.getAcid(before.start).secondary_structure;
                for (let acid of protein.getAcids(location.start, before.start)) {
                    acid.secondary_structure = newSecs;
                }
                protein.invalidateLoopsGeometry(location.start, before.start)
            }

            if (before.end < location.end) {
                let newSecs = protein.getAcid(before.end).secondary_structure;
                for (let acid of protein.getAcids(before.end, location.end)) {
                    acid.secondary_structure = newSecs;
                }
                protein.invalidateLoopsGeometry(before.end, location.end)
            } else if (location.end < before.end) {
                for (let acid of protein.getAcids(location.end + 1, before.end)) {
                    acid.secondary_structure = 'C';
                }
                protein.invalidateLoopsGeometry(location.end, before.end)

            }
            loop.setPosition(location.start, location.end, protein);

        },

        addLoopFlexibility(state, { proteinName, loop_id, flexibility, crosscorelation }) {
            let loop = state.proteins[proteinName].loops.find(l => l.id == loop_id);
            loop.flexibility = flexibility;
            loop.crosscorelation = crosscorelation;

        },

        addProteinFlexibilityTypes(state, { proteinName, flexibility }) {
            state.proteins[proteinName]
                .availableBfactors
                .push(...flexibility.types);
            state.proteins[proteinName]
                .bfactorCorrelation = flexibility.correlation_matrix;
        },

        addProteinCcorsTypes(state, { proteinName, ccors }) {
            state.proteins[proteinName]
                .availableCcors
                .push(...ccors);

        },

        addProteinSecondaryStructureFlexibility(state, { proteinName, flexibility }) {
            let protein = state.proteins[proteinName];

            for (let flex of flexibility.values) {
                let bfactors = {};
                for (let flextype of flexibility.types) {
                    bfactors[flextype] = roundToTwo(parseFloat(flex[flextype]))
                }
                let segment = flexibility.segmentation.find(s => s.name == flex.segment);
                for (let acid_num = segment.from; acid_num <= segment.to; acid_num++) {
                    let acid = protein.getAcid(acid_num);
                    if (acid != null) {
                        acid.bfactors.secondary_structure = bfactors;
                    }
                }
            }

        },

        addProteinSecondaryStructureCrosscol(state, { proteinName, crosscol }) {
            let protein = state.proteins[proteinName];
            protein.secs_crosscol = {};
            for (let ccrosc_type in crosscol) {
                protein.secs_crosscol[ccrosc_type] = {};


                for (let secs in crosscol[ccrosc_type].segment_dictionary) {
                    let pairings = zipWith(
                        crosscol[ccrosc_type].segment_dictionary[secs],
                        crosscol[ccrosc_type].segment_names,
                        (value, secsId) => ({ secsId, value: roundToTwo(Number(value)) })
                    )
                    protein.secs_crosscol[ccrosc_type][secs] = keyBy(pairings, c => c.secsId);

                }
            }

        },
        addAcidsFlexibility(state, { proteinName, flexibility }) {
            let protein = state.proteins[proteinName];

            for (let flex of flexibility.values) {
                let acid = protein.getAcid(parseInt(flex.segment));
                let bfactors = {};

                for (let flextype of flexibility.types) {
                    bfactors[flextype] = roundToTwo(parseFloat(flex[flextype]))
                }
                if (acid != null) {
                    acid.bfactors.residue = bfactors;
                }
            }

        },

        addSuperloops(state, { proteinName, loops }) {
            state.proteins[proteinName].superloops = loops;
            state.proteins[proteinName].hasSuperloops = true;
        },
        setSecondaryStructures(state, { proteinName, from, to, type }) {
            let acids = state.proteins[proteinName].getAcids(from.pdb_number, to.pdb_number);
            for (let acid of acids) {
                acid.secondary_structure = type
            }
            state.proteins[proteinName].computedLoops = [];

        },

        deleteLoop(state, { proteinName, loop }) {
            let protein = state.proteins[proteinName];
            protein.computedLoops = protein.computedLoops.filter(l => l.id != loop.id);
            protein.archDBLoops = protein.archDBLoops.filter(l => l.id != loop.id);
            protein.customLoops = protein.customLoops.filter(l => l.id != loop.id);

        },

        renewSegmentation(state, { proteinName }) {
            let protein = state.proteins[proteinName];
            protein.segmentation = new Segmentation(protein);
        },

        setProteinSelection(state, { proteinName, selection }) {
            state.proteins[proteinName].selection = selection;
        }

    },
    actions: {

        deleteLoops(context, { proteinName }) {
            let protein = context.getters.getProtein(proteinName);
            for (let loop of protein.loops) {
                context.dispatch("deleteLoop", { proteinName, loop })
            }
        },
        purgeNoncustomLoops(context, { proteinName }) {
            let protein = context.getters.getProtein(proteinName);
            for (let loop of [...protein.archDBLoops, ...protein.computedLoops]) {
                context.dispatch("deleteLoop", { proteinName, loop })
            }
        },
        deleteLoop(context, { proteinName, loop }) {
            context.commit("deleteLoop", { proteinName, loop });
            context.commit("deleteInterestingLoop", loop);
            context.dispatch("removeLoopDetailByLoop", loop);
        },


        setSecondaryStructures(context, { proteinName, from, to, type }) {

            for (let loop of context.getters.getProtein(proteinName).computedLoops) {
                context.commit("deleteInterestingLoop", loop);
            }
            context.commit("setSecondaryStructures", { proteinName, from, to, type });
        },

        select(context, selection) {
            context.commit("select", selection);
            if (context.getters.getCurrentSelection.isSelected()) {
                let currentSelection = context.state.currentSelection;
                context.commit("completeSelection");
                return currentSelection;
            }
            return null;
        },

        addCustomLoop(context, { proteinName, region }) {
            let protein = context.getters.getProtein(proteinName);

            while (protein.allLoops.map(l => l.id).includes(region.name)) {
                region.name = region.name + "'";
            }
            let from = region.from.pdb_number;
            let to = region.to.pdb_number;

            while (protein.getAcid(from).simple_secondary_structure == "C"
                && from < to) {
                from += 1;
            }

            while (protein.getAcid(to).simple_secondary_structure == "C"
                && from < to) {
                to -= 1;
            }

            let geometry = {
                distance: region.distance,
                rho: region.rho,
                delta: region.delta,
                theta: region.theta
            }


            for (let computedLoop of protein.computedLoops.concat(protein.computedSuperloops)) {
                if (computedLoop.start == from && computedLoop.end == to) {
                    geometry = {
                        distance: computedLoop.distance,
                        rho: computedLoop.rho,
                        delta: computedLoop.delta,
                        theta: computedLoop.theta
                    }
                }
            }




            let newLoop = new Loop(region.name,
                from,
                to,
                to - from,
                protein.getAcids(from, to),
                protein.name,
                protein.chain,
                protein,
                "Custom",
                geometry);
            context.commit("addCustomLoop", { proteinName, loop: newLoop })
            context.dispatch("makeInterestingLoop", { rawLoop: newLoop })
        },

        makeSelectionFromAcids(context, { protein, acids }) {
            context.commit("cancelSelection");
            context.commit("select", {
                protein,
                acid: acids.first,
            })
            context.commit("select", {
                protein,
                acid: acids.last,
            })
            context.commit("completeSelection");
        },

        removeProtein(context, proteinName) {
            context.commit("removeProtein", proteinName);

            //TODO: make remove loop/selection by protein name mutations for easy removal instead of this BS
            //TODO: actually maybe resruct the single source of truth - I need a protein structure that contains it's props
            //instead of multiple objects. Or at least do some kind of map.
            for (let detail of context.getters.getLoopDetails) {
                if (detail.protein.name == proteinName) {
                    context.commit("removeLoopDetail", detail.protein.name + detail.loop.id);
                }
            }

            context.dispatch("deleteInterestingLoops", proteinName);

            if (context.getters.getThreeDProtein?.name == proteinName) {
                context.commit('setThreeDProtein', { protein: null })
            }
            // for(let selection of Object.values(context.getters.getLoopDetails)) {
            //     if (selection.protein.name == proteinName) {
            //         context.commit("remove")
            //     }
            // }
        },

        _addProtein(context, { proteinName, chain, role }) {
            let api = context.getters.getApi;
            chain = chain.toUpperCase();
            let pdbPromise =
                api.get(`/downloadpdbs/${proteinName}/${chain}`)
                    .catch(err => {
                        context.dispatch("setErrorMessage", { error: err });
                        throw err;
                    });

            let dsspPromise = api.get(`/dssp/${proteinName}`)
                .catch(err => {
                    context.dispatch("setErrorMessage", { error: err })
                    throw err;
                });


            context.dispatch("addAwaitable", {
                id: _.uniqueId("awaiting"),
                promise: dsspPromise,
                description: "Computing DSSP assingment"
            });

            context.dispatch("addAwaitable", {
                id: _.uniqueId("awaiting"),
                promise: pdbPromise,
                description: "Getting PDB"
            });

            return Promise.all([pdbPromise, dsspPromise])
                .then(function ([pdb, dssp]) {
                    debugger;

                    let newProtein = new Protein(proteinName, pdb.data, dssp.data, chain, role);
                    context.commit("addProtein", newProtein);
                    context.commit(role == "scaffold" ? "setScaffold" : "setInsert", newProtein.name);
                    return newProtein
                }).then(function (protein) {
                    return context.commit("setThreeDProtein", { protein: protein, flexibility: null });
                })
                .catch((err) => {
                    context.dispatch("setErrorMessage", { error: err });
                    throw err;
                })
        },

        addRCBSProtein(context, { proteinName, config, role }) {
            let request = context.getters.getApi.post(`/rcsb`, { "protein": proteinName.trim() })
                .then(() => context.dispatch("_addProtein", { proteinName, ...config, role }))
                .catch(err => {
                    context.dispatch("setErrorMessage", { error: err })
                    throw err;
                });
            context.dispatch("addAwaitable", {
                id: _.uniqueId("awaiting"),
                promise: request,
                description: "Downloading from RCSB"
            });
            return request;
        },

        uploadProtein(context, { name, file, chain, role }) {
            let request = toBase64(file)
                .then(encodedFile => context.getters.getApi.post(`/userpdb`, {
                    pdbname: name,
                    pdbfile: encodedFile
                }))
                .catch(err => {
                    context.dispatch("setErrorMessage", { error: err });
                    throw err;
                })
                .then(() => context.dispatch("_addProtein", {
                    proteinName: name,
                    chain,
                    computeLoops: false,
                    role

                }));
            context.dispatch("addAwaitable", {
                id: _.uniqueId("awaiting"),
                promise: request,
                description: "Uploading custom PDB file"
            });
            return request
        },

        async addCorrelatedBFactors(context, { protein, bfactors, files }) {
            context.commit("renewSegmentation", { proteinName: protein.name });
            let flexibility_files = await Promise.all(files.map(f => toBase64(f)));
            let flexibility_computation = Promise.all(
                [
                    (() => {
                        let residue_proimise = context.getters.getApi.post(`/flexibility/submit`, {
                            pdb: protein.name,
                            chain: protein.chain,
                            types: bfactors,
                            segmentation: null,
                            file: flexibility_files,
                            crosscorelation: false,
                            bfactor_aggregation_type: 'residue',
                        });
                        return residue_proimise.then(
                            (response) => wait_for_computation_async(context.getters.getApi, "flexibility", 5000, response.data.taskId)
                        )
                    })(),
                    (() => {
                        let loops_promise = context.getters.getApi.post(`/flexibility/submit`, {
                            pdb: protein.name,
                            chain: protein.chain,
                            types: bfactors,
                            segmentation: protein.segmentation
                                .segments
                                .filter(segment => segment.type == SegmentationType.loop),
                            file: flexibility_files,
                            bfactor_aggregation_type: 'loop',
                            crosscorelation: true
                        });
                        return loops_promise.then(
                            (response) => wait_for_computation_async(context.getters.getApi, "flexibility", 5000, response.data.taskId)
                        )
                    })(),
                    (() => {
                        let secs_promise = context.getters.getApi.post(`/flexibility/submit`, {
                            pdb: protein.name,
                            chain: protein.chain,
                            types: bfactors,
                            segmentation: protein.segmentation
                                .segments
                                .filter(segment => segment.type == SegmentationType.secondary_structure),
                            file: flexibility_files,
                            bfactor_aggregation_type: 'secs'

                        });
                        return secs_promise.then(
                            (response) => wait_for_computation_async(context.getters.getApi, "flexibility", 5000, response.data.taskId)
                        )
                    })()
                ]
            ).then(([residue_response, loops_response, secs_response]) => {
                if (!residue_response.data.canBeRetrieved) {
                    throw new Error("Computation of flexibility per-residue has failed.")
                }
                if (!loops_response.data.canBeRetrieved) {
                    throw new Error("Computation of flexibility per-loop has failed.")
                }
                if (!secs_response.data.canBeRetrieved) {
                    throw new Error("Computation of flexibility per-secs has failed.")
                }

                return Promise.all([
                    context.getters.getApi.get(`/flexibility/${residue_response.data.taskId}/results`),
                    context.getters.getApi.get(`/flexibility/${loops_response.data.taskId}/results`),
                    context.getters.getApi.get(`/flexibility/${secs_response.data.taskId}/results`),
                ])

            }).then((([residues, loops, secs]) => {
                //residues
                debugger;
                context.commit("addAcidsFlexibility", {
                    proteinName: protein.name,
                    flexibility: residues.data
                })
                //loops
                for (let flex of loops.data.values) {
                    let flexibility_dict = {}
                    for (let flextype of loops.data.types) {
                        flexibility_dict[flextype] = roundToTwo(parseFloat(flex[flextype]))
                    }

                    let segment = loops.data.segmentation.find(s => s.name == flex.segment);

                    let crosscorelation_types = Object.keys(loops.data.cross_correlation);
                    let ccors = {}
                    for (let type of crosscorelation_types) {
                        let ccors_dictionary = {};
                        let ccors_segments = loops.data.cross_correlation[type].segment_names;
                        ccors_dictionary = zipWith(
                            loops.data.cross_correlation[type].segment_dictionary[segment.name],
                            ccors_segments,
                            (value, loopId) => ({ loopId, value: roundToTwo(parseFloat(value)) })
                        )
                        ccors[type] = keyBy(ccors_dictionary, c => c.loopId);
                    }
                    context.commit("addLoopFlexibility", {
                        proteinName: protein.name,
                        loop_id: segment.name,
                        flexibility: flexibility_dict,
                        crosscorelation: ccors
                    })
                    context.commit("addProteinCcorsTypes", {
                        proteinName: protein.name,
                        ccors: crosscorelation_types
                    })
                }

                context.commit("addProteinFlexibilityTypes", {
                    proteinName: protein.name,
                    flexibility: loops.data
                })
                //Secs
                context.commit("addProteinSecondaryStructureCrosscol", {
                    proteinName: protein.name,
                    crosscol: secs.data.cross_correlation
                });
                context.commit("addProteinSecondaryStructureFlexibility", {
                    proteinName: protein.name,
                    flexibility: secs.data
                })
            })).catch(err => {
                context.dispatch("setErrorMessage", { error: err });
                throw err;
            });

            context.dispatch("addAwaitable", {
                id: _.uniqueId("awaiting"),
                promise: flexibility_computation,
                description: `Computing flexibility for ${protein.displayName}`
            });
            return flexibility_computation;
        },

        addLoops(context, { proteinName, geometries, source }) {
            let protein = context.getters.getProtein(proteinName)
            let displayName = protein.displayName;
            for (let [index, rawLoop] of _.sortBy(geometries.loops, l => l.start).entries()) {
                let loop = new Loop(
                    `${displayName}_${rawLoop.chain}_${index + 1}`,
                    rawLoop.start,
                    rawLoop.end,
                    protein.getAcids(rawLoop.start, rawLoop.end),
                    rawLoop.length,
                    rawLoop.pdb,
                    rawLoop.chain,
                    protein,
                    source,
                    {
                        distance: rawLoop.distance,
                        delta: rawLoop.delta,
                        rho: rawLoop.rho,
                        theta: rawLoop.theta
                    }
                )
                context.commit("addComuptedLoop", {
                    proteinName,
                    loop
                })
                context.dispatch("makeInterestingLoop", { rawLoop: loop })
            }

            for (let rawLoop of geometries.superloops) {
                let loop = new Loop(
                    `${displayName}_${rawLoop.chain}_${rawLoop.start}_${rawLoop.end}`,
                    rawLoop.start,
                    rawLoop.end,
                    protein.getAcids(rawLoop.start, rawLoop.end),
                    rawLoop.length,
                    rawLoop.pdb,
                    rawLoop.chain,
                    protein,
                    source,
                    {
                        distance: rawLoop.distance,
                        delta: rawLoop.delta,
                        rho: rawLoop.rho,
                        theta: rawLoop.theta
                    }
                )
                context.commit("addComputedSuperLoop", {
                    proteinName,
                    loop
                })

            }

        },

        computeLoops(context, proteinName) {
            console.log("starting loop comp...")
            let protein = context.getters.getProtein(proteinName);
            let request = context.getters.getApi.post(`/loops/submit`, {
                pdb: proteinName,
                chain: protein.chain,
                segmentation: new Segmentation(protein)
                    .segments
                    .filter(segment => segment.type == SegmentationType.secondary_structure)
            }).then((response) => {   //check until finished
                console.log(response)
                return Promise.all([wait_for_computation_async(context.getters.getApi, "loops", 5000, response.data.taskId), response.data.taskId])
            }).then(([response, taskId]) => {
                console.log(response)
                if (!response.data.canBeRetrieved) {
                    throw new Error("Loops computation failed with unknown error.");
                }
                return context.getters.getApi.get(`/loops/${taskId}/results`)
            }).then((response) => { //set
                context.dispatch("purgeNoncustomLoops", { proteinName: protein.name });
                context.dispatch("addLoops", { proteinName, geometries: response.data, source: "Computed" })
            }).catch(err => {
                context.dispatch("setErrorMessage", { error: err });
                throw err;
            });
            context.dispatch("addAwaitable", {
                id: _.uniqueId("awaiting"),
                promise: request,
                description: `Computing loops for ${protein.displayName}`
            });

        },

        zoomToLoop(context, { protein, loop }) {
            context.commit("setThreeDProtein", { protein: protein, flexibility: null });
            context.commit("zoomToLoop", loop);
        },
        removeRole(context, role) {
            if (role == "scaffold") {
                let proteinName = context.getters.getScaffold?.name;
                context.dispatch("removeProtein", proteinName);
                context.commit("setScaffold", null);
            } else {
                let proteinName = context.getters?.getInsert;
                context.dispatch("removeProtein", proteinName);
                context.commit("setInsert", null);
            }
        },

    },

}

