# -*- coding: utf-8 -*-
# SyConn - Synaptic connectivity inference toolkit
#
# Copyright (c) 2016 - now
# Max-Planck-Institute of Neurobiology, Munich, Germany
# Authors: Philipp Schubert, Joergen Kornfeld

# TODO: move code to syconn/gate/ but keep executable script here
import copy
import time
import numpy as np
from flask import Flask
import json
from multiprocessing import cpu_count
import argparse
import os

from syconn.reps import super_segmentation as ss
from syconn.reps import connectivity_helper as conn
from syconn import global_params
from syconn.handler.logger import log_main as log_gate
from syconn.handler.prediction import int2str_converter
from syconn.reps.segmentation import SegmentationDataset

app = Flask(__name__)
sg_state = None


@app.route('/ssv_skeleton/<ssv_id>', methods=['GET'])
def route_ssv_skeleton(ssv_id):
    d = sg_state.backend.ssv_skeleton(ssv_id)
    start = time.time()
    ret = json.dumps(d, cls=MyEncoder)
    log_gate.debug("JSON dump: {}".format(time.time() - start))
    return ret


@app.route('/ssv_mesh/<ssv_id>', methods=['GET'])
def route_ssv_mesh(ssv_id):
    d = sg_state.backend.ssv_mesh(ssv_id)
    start = time.time()
    ret = json.dumps(d, cls=MyEncoder)
    log_gate.debug("JSON dump: {}".format(time.time() - start))
    return ret


@app.route('/ssv_ind/<ssv_id>', methods=['GET'])
def route_ssv_ind(ssv_id):
    d = sg_state.backend.ssv_ind(ssv_id)
    return d


@app.route('/ssv_vert/<ssv_id>', methods=['GET'])
def route_ssv_vert(ssv_id):
    d = sg_state.backend.ssv_vert(ssv_id)
    return d


@app.route('/ssv_norm/<ssv_id>', methods=['GET'])
def route_ssv_norm(ssv_id):
    d = sg_state.backend.ssv_norm(ssv_id)
    return d


@app.route('/ssv_obj_vert/<ssv_id>/<obj_type>', methods=['GET'])
def ssv_obj_vert(ssv_id, obj_type):
    d = sg_state.backend.ssv_obj_vert(ssv_id, obj_type)
    return d


@app.route('/ssv_obj_ind/<ssv_id>/<obj_type>', methods=['GET'])
def ssv_obj_ind(ssv_id, obj_type):
    d = sg_state.backend.ssv_obj_ind(ssv_id, obj_type)
    return d


@app.route('/ssv_obj_norm/<ssv_id>/<obj_type>', methods=['GET'])
def ssv_obj_norm(ssv_id, obj_type):
    d = sg_state.backend.ssv_obj_norm(ssv_id, obj_type)
    return d


@app.route('/ssv_list', methods=['GET'])
def route_ssv_list():
    return json.dumps(sg_state.backend.ssv_list(), cls=MyEncoder)


@app.route('/pull_so_attr/<so_id>/<so_type>/<attr_key>', methods=['GET'])
def pull_so_attr(so_id, so_type, attr_key):
    return json.dumps(sg_state.backend.pull_so_attr(so_id, so_type, attr_key),
                      cls=MyEncoder)


@app.route('/push_so_attr/<so_id>/<so_type>/<attr_key>/<attr_value>', methods=['GET'])
def push_so_attr(so_id, so_type, attr_key, attr_value):
    return json.dumps(sg_state.backend.push_so_attr(so_id, so_type, attr_key, attr_value),
                      cls=MyEncoder)


@app.route('/svs_of_ssv/<ssv_id>', methods=['GET'])
def route_svs_of_ssv(ssv_id):
    return json.dumps(sg_state.backend.svs_of_ssv(ssv_id), cls=MyEncoder)


@app.route('/ssv_of_sv/<sv_id>', methods=['GET'])
def route_ssv_of_sv(sv_id):
    return json.dumps(sg_state.backend.ssv_of_sv(sv_id), cls=MyEncoder)


@app.route('/ct_of_ssv/<ssv_id>', methods=['GET'])
def route_ct_of_sv(ssv_id):
    return json.dumps(sg_state.backend.ct_of_ssv(ssv_id), cls=MyEncoder)


@app.route('/all_syn_meta/<params>', methods=['GET'])
def route_all_syn_meta(params):
    return json.dumps(sg_state.backend.all_syn_meta_data(json.loads(params)), cls=MyEncoder)


@app.route("/", methods=['GET'])
def route_hello():
    return json.dumps({'Welcome to': 'SyConnGate'}, cls=MyEncoder)


class SyConnBackend(object):
    def __init__(self, syconn_path: str = '', logger=None, synthresh=0.5,
                 axodend_only=True):
        """
        Initializes a SyConn backend for operation.
        This includes in-memory initialization of the
        most important caches. Currently, SyConn Gate
        does not support backend data changes and the server needs
        to restart for changes to be valid. If the backend data
        is changed while the server is running, old content
        might be served.
        All backend functions must return dicts.

        Args:
            syconn_path:
            logger:
            synthresh: All synapses below `synthresh` will be excluded.
            axodend_only: If True, only axo-dendritic synapses will be loaded.
        """
        self.logger = logger
        self.logger.info('Initializing SyConn backend')

        self.ssd = ss.SuperSegmentationDataset(syconn_path,
                                               # deal with concurrency issues
                                               sso_locking=False,
                                               sso_caching=True)

        self.logger.info('SuperSegmentation dataset initialized.')

        self.sds = dict(syn_ssv=SegmentationDataset(working_dir=syconn_path,
                                                    obj_type='syn_ssv'))
        self.nb_cpus = cpu_count()
        self.synthresh = synthresh
        self.axodend_only = axodend_only
        # flat array representation of all synapses
        self.conn_dict = conn.load_cached_data_dict()
        self.logger.info(f'In-memory cache of {len(self.sds["syn_ssv"].ids)} synapses initialized.')
        # directed networkx graph of connectivity
        self.conn_graph = conn.connectivity_to_nx_graph(self.conn_dict)
        self.logger.info('Connectivity graph initialized.')

    def ssv_mesh(self, ssv_id):
        """
        Get mesh for ssv_id.
        Args:
            ss_vid (int): id of ssv
        Returns:
            dict:
        """
        start = time.time()
        self.logger.info('Loading ssv mesh {}'.format(ssv_id))
        ssv = self.ssd.get_super_segmentation_object(int(ssv_id))
        ssv.nb_cpus = self.nb_cpus
        ssv.load_attr_dict()
        mesh = ssv.load_mesh("sv")
        mesh = {'vertices': mesh[1],
                'indices': mesh[0],
                'normals': []}
        dtime = time.time() - start
        self.logger.info('Got ssv mesh {} after {:.2f}'.format(ssv_id, dtime))
        return mesh

    def ssv_ind(self, ssv_id):
        """
        Args:
            ssv_id (int):
        Returns:
            dict:
        """
        start = time.time()
        self.logger.info('Loading {} ssv mesh indices'.format(ssv_id))
        ssv = self.ssd.get_super_segmentation_object(int(ssv_id))
        ssv.nb_cpus = self.nb_cpus
        ssv.load_attr_dict()
        mesh = ssv.load_mesh("sv")
        dtime = time.time() - start
        self.logger.info('Got ssv {} mesh indices after'
                         ' {:.2f}'.format(ssv_id, dtime))
        return {'ind': mesh[0].tolist()}

    def ssv_vert(self, ssv_id):
        """
        Get mesh vertices for ssv_id.
        Args:
            ssv_id (int):
        Returns:
            dict:
        """
        start = time.time()
        self.logger.info('Loading ssv {} mesh vertices'.format(ssv_id))
        ssv = self.ssd.get_super_segmentation_object(int(ssv_id))
        ssv.nb_cpus = self.nb_cpus
        ssv.load_attr_dict()
        mesh = ssv.load_mesh("sv")
        dtime = time.time() - start
        self.logger.info('Got ssv {} mesh vertices after'
                         ' {:.2f}'.format(ssv_id, dtime))
        return {'vert': mesh[1].tolist()}

    def ssv_skeleton(self, ssv_id):
        """
        Get mesh vertices for ssv_id.
        Args:
            ssv_id (int):
        Returns:
            dict:
        """
        self.logger.info('Loading ssv skeleton {}'.format(ssv_id))
        ssv = self.ssd.get_super_segmentation_object(int(ssv_id))
        ssv.nb_cpus = self.nb_cpus
        ssv.load_skeleton()
        skeleton = ssv.skeleton
        if skeleton is None:
            return {}
        skel_attr = ["nodes", "edges", "diameters"]
        pred_key_ax = "{}_avg{}".format(ssv.config['compartments']['view_properties_semsegax']['semseg_key'],
                                        ssv.config['compartments']['dist_axoness_averaging'])
        pred_key_sp = ssv.config['spines']['semseg2mesh_spines']['semseg_key']
        keys = [
                global_params.config['compartments']['view_properties_semsegax']['semseg_key'],
                pred_key_ax,
                pred_key_ax + '_comp_maj',
                pred_key_sp,
                'myelin_avg10000',  # TODO: use global_params.py value !
                'myelin']  # TODO: use global_params.py value !
        for k in keys:
            if k in skeleton:
                skel_attr.append(k)
                if type(skeleton[k]) is list:
                    skeleton[k] = np.array(skeleton[k])
            else:
                log_gate.warning("Couldn't find requested key in "
                                 "skeleton '{}'. Existing keys: {}".format(k, skeleton.keys()))
        return {k: skeleton[k].flatten().tolist() for k in
                skel_attr}

    def ssv_norm(self, ssv_id):
        """
        Get mesh normals for ssv_id.
        Args:
            ssv_id (int):
        Returns:
            dict:
        """
        # not needed for K
        return {'norm': []}
        start = time.time()
        self.logger.info('Loading ssv {} mesh normals'.format(ssv_id))
        ssv = self.ssd.get_super_segmentation_object(int(ssv_id))
        ssv.nb_cpus = self.nb_cpus
        ssv.load_attr_dict()
        mesh = ssv.mesh("sv")
        dtime = time.time() - start
        self.logger.info('Got ssv {} mesh normals after'
                         ' {:.2f}'.format(ssv_id, dtime))
        if len(mesh) == 2:
            return {'norm': []}
        return {'norm': mesh[2].tolist()}

    def ssv_obj_ind(self, ssv_id, obj_type):
        """
        Get mesh indices of a specific obj type for ssv_id.
        Args:
            ssv_id (int):
            obj_type (str):
        Returns:
            dict:
        """
        start = time.time()
        self.logger.info('Loading ssv {} {} mesh indices'
                         ''.format(ssv_id, obj_type))
        ssv = self.ssd.get_super_segmentation_object(int(ssv_id))
        ssv.nb_cpus = self.nb_cpus
        ssv.load_attr_dict()
        if obj_type == "sj":
            try:
                obj_type = "syn_ssv"
                _ = ssv.attr_dict[obj_type]  # try to query mapped syn_ssv objects
                log_gate.debug("Loading '{}' objects instead of 'sj' for SSV "
                               "{}.".format(obj_type, ssv_id))
            except KeyError:
                obj_type = "sj"
        # Now assumes all object meshes do already exist
        mesh = ssv.load_mesh(obj_type)
        dtime = time.time() - start
        self.logger.info('Got ssv {} {} mesh indices after'
                         ' {:.2f}'.format(ssv_id, obj_type, dtime))
        return {'ind': mesh[0].tolist()}

    def ssv_obj_vert(self, ssv_id, obj_type):
        """
        Get mesh vertices  of a specific obj type for ssv_id.
        Args:
            ssv_id (int):
            obj_type (str):
        Returns:
            dict:
        """
        start = time.time()
        self.logger.info('Loading ssv {} {} mesh vertices'
                         ''.format(ssv_id, obj_type))
        ssv = self.ssd.get_super_segmentation_object(int(ssv_id))
        ssv.nb_cpus = self.nb_cpus
        ssv.load_attr_dict()
        if obj_type == "sj":
            try:
                obj_type = "syn_ssv"
                _ = ssv.attr_dict[obj_type]  # try to query mapped syn_ssv objects
                log_gate.debug("Loading '{}' objects instead of 'sj' for SSV "
                               "{}.".format(obj_type, ssv_id))
            except KeyError:
                obj_type = "sj"
        # if not existent, create mesh
        mesh = ssv.load_mesh(obj_type)
        dtime = time.time() - start
        self.logger.info('Got ssv {} {} mesh vertices after'
                         ' {:.2f}'.format(ssv_id, obj_type, dtime))
        return {'vert': mesh[1].tolist()}

    def ssv_obj_norm(self, ssv_id, obj_type):
        """
        Get mesh normals of a specific obj type for ssv_id.
        Args:
            ssv_id (int):
            obj_type (str):
        Returns:
            dict:
        """
        return {'norm': []}
        start = time.time()
        self.logger.info('Loading ssv {} {} mesh normals'
                         ''.format(ssv_id, obj_type))
        ssv = self.ssd.get_super_segmentation_object(int(ssv_id))
        ssv.nb_cpus = self.nb_cpus
        ssv.load_attr_dict()
        if obj_type == "sj":
            try:
                obj_type = "syn_ssv"
                _ = ssv.attr_dict[obj_type]  # try to query mapped syn_ssv objects
                log_gate.debug("Loading '{}' objects instead of 'sj' for SSV "
                               "{}.".format(obj_type, ssv_id))
            except KeyError:
                obj_type = "sj"
        # if not existent, create mesh
        mesh = ssv.load_mesh(obj_type)
        dtime = time.time() - start
        self.logger.info('Got ssv {} {} mesh normals after'
                         ' {:.2f}'.format(ssv_id, obj_type, dtime))
        if len(mesh) == 2:
            return {'norm': []}
        return {'norm': mesh[2].tolist()}

    def ssv_list(self):
        """
        Returns all ssvs in dataset.
        Return:
            dict:
        """
        return {'ssvs': list(self.ssd.ssv_ids)}

    def ssv_of_sv(self, sv_id):
        """
        Returns the ssv for a given sv_id.
        Args:
            sv_id: 
        Returns:
        """
        sv_id = int(sv_id)
        try:
            ssv_id = self.ssd.sv2ssv_ids([int(sv_id)])[sv_id]
        except KeyError:
            # sv_id is not in sv2ssv_ids output (dict) -> ssv ID is not part of ssd.ssv_ids
            ssv_id = 0
        return {'ssv': ssv_id}

    def ct_of_ssv(self, ssv_id):
        """
        Returns the CT for a given SSV ID.
        Args:
            sv_id: 
        Returns:
        """
        ssv = self.ssd.get_super_segmentation_object(int(ssv_id))
        ssv.nb_cpus = self.nb_cpus
        ssv.load_attr_dict()
        label = ""
        certainty = ""
        # TODO: generalize!
        gt_type = 'ctgt_v2'
        if 'j0251' in self.ssd.working_dir:
            gt_type = 'ctgt_j0251_v2'
        if "celltype_cnn_e3_probas" in ssv.attr_dict:  # new prediction
            l = ssv.attr_dict["celltype_cnn_e3"]
            # ct_label_dc = {0: "EA", 1: "MSN", 2: "GP", 3: "INT"}
            # label = ct_label_dc[l]
            label = int2str_converter(l, gt_type=gt_type)
            certainty = ssv.certainty_celltype()
        elif "celltype_cnn" in ssv.attr_dict:
            ct_label_dc = {0: "EA", 1: "MSN", 2: "GP", 3: "INT"}
            l = ssv.attr_dict["celltype_cnn"]
            label = ct_label_dc[l]
            certainty = 'nan'
        else:
            log_gate.warning("Celltype prediction not present in attribute "
                             "dict of SSV {} at {}.".format(ssv_id, ssv.attr_dict_path))
        return {'ct': label, 'certainty': certainty}

    def svs_of_ssv(self, ssv_id):
        """
        Returns all sv ids for a ssv_id.
        Returns:
            dict:
        """
        ssv = self.ssd.get_super_segmentation_object(int(ssv_id))
        ssv.nb_cpus = self.nb_cpus
        ssv.load_attr_dict()
        return {'svs': ssv.sv_ids.tolist()}

    def all_syn_meta_data(self, params):
        """
        Returns all synapse meta data. This works only well for fast
        connections and less than 1e6 synapses or so. `synthresh` is transmitted
        with a scaling factor of 1000 (-> maximum precision is 4 digits).
        Returns:
        """
        synthresh = params['synthresh']
        axodend_only = params['axodend_only']
        if (synthresh != self.synthresh) or (axodend_only != self.axodend_only):
            self.synthresh = synthresh
            # flat array representation of all synapses
            self.conn_dict = conn.load_cached_data_dict(thresh_syn_prob=synthresh,
                                                        axodend_only=axodend_only)
            self.logger.info(f'In memory cache of synapses re-initialized with '
                             f'threshold {synthresh} and "axodend_only={axodend_only}".')
            # directed networkx graph of connectivity
            self.conn_graph = conn.connectivity_to_nx_graph(self.conn_dict)
            self.logger.info('Connectivity graph re-initialized.')
        all_syn_meta_dict = copy.copy(self.conn_dict)

        # the self.conn_dict is not json serializeble, due to the numpy arrays
        for k, v in all_syn_meta_dict.items():
            all_syn_meta_dict[k] = v.tolist()

        return all_syn_meta_dict

    def syn_objs_of_ssv_pre_post(self, ssv_id):
        """
        TODO: Requires adaptions of 'SyConnBackend' class
        Returns all synapse objs of a given ssv_id.
        Returns:
        """
        syns = dict()
        # not the most efficient approach, a cached map might be necessary for
        # many synapses
        idx = (ssv_id == self.syn_ssv_partner_0) | (
        ssv_id == self.syn_ssv_partner_1)

        syns['ids'] = self.syn_ids[idx]
        syns['sizes'] = self.syn_sizes[idx]
        syns['p0'] = self.syn_ssv_partner_0[idx]
        syns['p1'] = self.syn_ssv_partner_1[idx]
        return syns

    def syn_objs_of_ssv_post(self, ssv_id):
        """
        TODO: Requires adaptions of 'SyConnBackend' class
        Return the syn objs where this ssv_id is post synaptic,
        i.e. this ssv_id receives the synapse.

        Args:
            sv_id: 
        Returns:
        """
        syns = dict()
        # not the most efficient approach, a cached map might be necessary for
        # really many synapses
        idx = (ssv_id == self.syn_ssv_partner_0) | (
        ssv_id == self.syn_ssv_partner_1)

        # find the synapses where this ssv_id is post synaptic

        syns['ids'] = self.syn_ids[idx]
        syns['sizes'] = self.syn_sizes[idx]
        syns['p0'] = self.syn_ssv_partner_0[idx]
        syns['p1'] = self.syn_ssv_partner_1[idx]
        return syns

    def pull_so_attr(self, so_id, so_type, attr_key):
        """
        Generic attribute pull, return empty string if key did not exist. Could be optimized
        with the assumption that all attributes have been cached as numpy arrays.

        Args:
            so_id (int):
            so_type(str):
            attr_key (str):

        Returns:
            str:
        """
        if so_type not in self.sds:
            self.sds[so_type] = SegmentationDataset(obj_type=so_type)
        sd = self.sds[so_type]
        so = sd.get_segmentation_object(so_id)
        so.load_attr_dict()
        if attr_key not in so.attr_dict:
            return ''
        return so.attr_dict[attr_key]

    def push_so_attr(self, so_id, so_type, attr_key, attr_value):
        """
        Generic attribute pull, return empty string if key did not exist. Could be optimized
        with the assumption that all attributes have been cached as numpy arrays.

        Args:
            so_id (int):
            so_type(str):
            attr_key (str):
            attr_value

        Returns:
            bytes: Empty string of everything went well
        """
        
        if so_type not in self.sds:
            self.sds[so_type] = SegmentationDataset(obj_type=so_type)
        sd = self.sds[so_type]
        try:
            so = sd.get_segmentation_object(so_id)
            so.save_attributes([attr_key], [attr_value])
            return ""
        except Exception as e:
            return str(e)


class ServerState(object):
    def __init__(self, host=None, port=None):
        self.logger = log_gate
        self.host = host
        self.port = port

        self.logger.info('SyConn gate server starting up on working directory '
                         '"{}".'.format(global_params.wd))
        if not np.any(['syn_ssv' in name for name in os.listdir(global_params.config.working_dir)]):
            msg = 'Could not find synapse results in working directory ' \
                  f'{global_params.config.working_dir}.'
            self.logger.error(msg)
            raise RuntimeError(msg)
        try:
            self.backend = SyConnBackend(global_params.config.working_dir, logger=self.logger)
        except TypeError as e:
            msg = 'TypeError during server initialization. Make sure SyConn has' \
                  f' successfully finished the data set analysis. \nError: {str(e)}'
            self.logger.error(msg)
            raise TypeError(msg)
        self.logger.info('SyConn gate server running at {}, {}.'.format(
            self.host, self.port))
        return


class MyEncoder(json.JSONEncoder):
    """
    From https://stackoverflow.com/questions/27050108/convert-numpy-type-to-python/27050186#27050186
    """
    def default(self, obj):
        if isinstance(obj, np.integer):
            return int(obj)
        elif isinstance(obj, np.floating):
            return float(obj)
        elif isinstance(obj, np.ndarray):
            return obj.tolist()
        else:
            return super(MyEncoder, self).default(obj)


def main():
    """
    Alternative way of running the server is currently:
    export FLASK_APP=server.py
    flask run --host=0.0.0.0 --port=10001 --debugger

    OR

    FLASK_APP=server.py FLASK_DEBUG=1 flask run --host=0.0.0.0 --port 10001

    """
    global sg_state
    parser = argparse.ArgumentParser(description='SyConn Gate')
    parser.add_argument('--working_dir', type=str,
                        default=global_params.config.working_dir,
                        help='Working directory of SyConn')
    parser.add_argument('--port', type=int, default=10001,
                        help='Port to connect to SyConn Gate')
    parser.add_argument('--host', type=str, default='0.0.0.0',
                        help='IP address to SyConn Gate')
    args = parser.parse_args()
    if args.working_dir is None:
        raise RuntimeError('SyConn working directory not set. Please use the '
                           '"--working_dir" argument or set "syconn_wd" in your'
                           ' shell environment.')
    server_wd = os.path.expanduser(args.working_dir)
    server_port = args.port
    server_host = args.host
    global_params.wd = server_wd

    sg_state = ServerState(server_host, server_port)

    # context = ('cert.crt', 'key.key') enable later
    app.run(host=server_host,  # do not run this on a non-firewalled machine!
            port=server_port,  # ssl_context=context,
            threaded=False, debug=False, use_reloader=False)


if __name__ == '__main__':
    main()
