# -*- coding: utf-8 -*-
"""
/***************************************************************************
 Simulator
                                 A QGIS plugin
 This plugin similaté the agent's movements.
 Generated by Plugin Builder: http://g-sherman.github.io/Qgis-Plugin-Builder/
                              -------------------
        begin                : 2021-07-10
        git sha              : $Format:%H$
        copyright            : (C) 2021 by HCMUT
        email                : duc.nguyenquang@hcmut.edu.vn
 ***************************************************************************/

/***************************************************************************
 *                                                                         *
 *   This program is free software; you can redistribute it and/or modify  *
 *   it under the terms of the GNU General Public License as published by  *
 *   the Free Software Foundation; either version 2 of the License, or     *
 *   (at your option) any later version.                                   *
 *                                                                         *
 ***************************************************************************/
"""
from qgis.PyQt.QtCore import *
from qgis.PyQt.QtGui import QIcon
from qgis.PyQt.QtWidgets import QAction, QFileDialog
from qgis.core import *

# Initialize Qt resources from file resources.py
from .resources import *
# Import the code for the dialog
from .simulator_dialog import SimulatorDialog
import os.path
import numpy as np
import random
import time
import functools
from sklearn.neighbors import KNeighborsClassifier

from .status import Status
from .agent import Agent

class Simulator:
    """QGIS Plugin Implementation."""

    def __init__(self, iface):
        """Constructor.

        :param iface: An interface instance that will be passed to this class
            which provides the hook by which you can manipulate the QGIS
            application at run time.
        :type iface: QgsInterface
        """
        # Save reference to the QGIS interface
        self.iface = iface
        # initialize plugin directory
        self.plugin_dir = os.path.dirname(__file__)
        # initialize locale
        locale = QSettings().value('locale/userLocale')[0:2]
        locale_path = os.path.join(
            self.plugin_dir,
            'i18n',
            'Simulator_{}.qm'.format(locale))

        if os.path.exists(locale_path):
            self.translator = QTranslator()
            self.translator.load(locale_path)
            QCoreApplication.installTranslator(self.translator)

        # Declare instance attributes
        self.actions = []
        self.menu = self.tr(u'&Simulator')

        # Check if plugin was started the first time in current QGIS session
        # Must be set in initGui() to survive plugin reloads
        self.first_start = None
        self.timer = None
        self.baseLayer = None
        self.step = 0
        self.run_steps = 0

        self.beta = 0
        self.gamma = 0
        self.immune = 0
        self.mu = 0
        self.step_mean = 80 # m
        self.step_std = 10 # m
        self.infectious_radius = 0 # m
        self.default_n_neighbor = 10

        self._susceptible = {}
        self._infectious = {}
        self._recovered = {}
        self._deceased = {}

        self._mapping_ward_id = {}
        self._mapping_agent_id = {}
        self._llist_adjacent_ward = {}

        self.knn = None

        self._sus_layer = QgsVectorLayer("MultiPoint?crs=epsg:9210", "Susceptible", "memory")
        pr_sus = self._sus_layer.dataProvider()
        pr_sus.addAttributes([QgsField("id", QVariant.String)])
        self._sus_layer.updateFields()

        self._inf_layer = QgsVectorLayer("MultiPoint?crs=epsg:9210", "Infectious", "memory")
        pr_inf = self._inf_layer.dataProvider()
        pr_inf.addAttributes([QgsField("id", QVariant.String)])
        self._inf_layer.updateFields()

        self._rec_layer = QgsVectorLayer("MultiPoint?crs=epsg:9210", "Recovered", "memory")
        pr_rec = self._rec_layer.dataProvider()
        pr_rec.addAttributes([QgsField("id", QVariant.String)])
        self._rec_layer.updateFields()

        self._dec_layer = QgsVectorLayer("MultiPoint?crs=epsg:9210", "Deceased", "memory")
        pr_dec = self._dec_layer.dataProvider()
        pr_dec.addAttributes([QgsField("id", QVariant.String)])
        self._dec_layer.updateFields()

    # noinspection PyMethodMayBeStatic
    def tr(self, message):
        """Get the translation for a string using Qt translation API.

        We implement this ourselves since we do not inherit QObject.

        :param message: String for translation.
        :type message: str, QString

        :returns: Translated version of message.
        :rtype: QString
        """
        # noinspection PyTypeChecker,PyArgumentList,PyCallByClass
        return QCoreApplication.translate('Simulator', message)


    def add_action(
        self,
        icon_path,
        text,
        callback,
        enabled_flag=True,
        add_to_menu=True,
        add_to_toolbar=True,
        status_tip=None,
        whats_this=None,
        parent=None):
        """Add a toolbar icon to the toolbar.

        :param icon_path: Path to the icon for this action. Can be a resource
            path (e.g. ':/plugins/foo/bar.png') or a normal file system path.
        :type icon_path: str

        :param text: Text that should be shown in menu items for this action.
        :type text: str

        :param callback: Function to be called when the action is triggered.
        :type callback: function

        :param enabled_flag: A flag indicating if the action should be enabled
            by default. Defaults to True.
        :type enabled_flag: bool

        :param add_to_menu: Flag indicating whether the action should also
            be added to the menu. Defaults to True.
        :type add_to_menu: bool

        :param add_to_toolbar: Flag indicating whether the action should also
            be added to the toolbar. Defaults to True.
        :type add_to_toolbar: bool

        :param status_tip: Optional text to show in a popup when mouse pointer
            hovers over the action.
        :type status_tip: str

        :param parent: Parent widget for the new action. Defaults None.
        :type parent: QWidget

        :param whats_this: Optional text to show in the status bar when the
            mouse pointer hovers over the action.

        :returns: The action that was created. Note that the action is also
            added to self.actions list.
        :rtype: QAction
        """

        icon = QIcon(icon_path)
        action = QAction(icon, text, parent)
        action.triggered.connect(callback)
        action.setEnabled(enabled_flag)

        if status_tip is not None:
            action.setStatusTip(status_tip)

        if whats_this is not None:
            action.setWhatsThis(whats_this)

        if add_to_toolbar:
            # Adds plugin icon to Plugins toolbar
            self.iface.addToolBarIcon(action)

        if add_to_menu:
            self.iface.addPluginToVectorMenu(
                self.menu,
                action)

        self.actions.append(action)

        return action

    def initGui(self):
        """Create the menu entries and toolbar icons inside the QGIS GUI."""

        icon_path = ':/plugins/simulator/icon.png'
        self.add_action(
            icon_path,
            text=self.tr(u'Initialize agents'),
            callback=self.run,
            parent=self.iface.mainWindow())

        self.add_action(
            icon_path,
            text=self.tr(u'Start simulation'),
            callback=self._run_simulation,
            parent=self.iface.mainWindow())

        # will be set False in run()
        self.first_start = True

    def unload(self):
        """Removes the plugin menu item and icon from QGIS GUI."""
        for action in self.actions:
            self.iface.removePluginVectorMenu(
                self.tr(u'&Simulator'),
                action)
            self.iface.removeToolBarIcon(action)

    def select_output_file(self):
        filename, _filter = QFileDialog.getOpenFileName(
        self.dlg, "Select input file ","", '*.csv')
        self.dlg.lineEdit.setText(filename)

    def run(self):
        """Run method that performs all the real work"""

        # Create the dialog with elements (after translation) and keep reference
        # Only create GUI ONCE in callback, so that it will only load when the plugin is started
        if self.first_start == True:
            self.first_start = False
            self.dlg = SimulatorDialog()
            self.dlg.pushButton.clicked.connect(self.select_output_file)

        # Fetch the currently loaded layers
        layers = QgsProject.instance().layerTreeRoot().children()
        # Clear the contents of the comboBox from previous runs
        self.dlg.comboBox.clear()
        # Populate the comboBox with names of all the loaded layers
        self.dlg.comboBox.addItems([layer.name() for layer in layers])

        # show the dialog
        self.dlg.show()
        # Run the dialog event loop
        result = self.dlg.exec_()
        # See if OK was pressed
        if result:
            # Do something useful here - delete the line containing pass and
            # substitute with your code.
            selectedLayerIndex = self.dlg.comboBox.currentIndex()
            selectedLayer = layers[selectedLayerIndex].layer()
            self.baseLayer = selectedLayer

            self.beta = self.dlg.box_beta.value()
            self.gamma = self.dlg.box_gamma.value()
            self.immune = self.dlg.box_immune.value()
            self.mu = self.dlg.box_mu.value()

            self._construct_adjacent_ward()

            # Number of day
            number_of_days = self.dlg.box_days.value()
            # number_of_days = 5/(24*60)

            # Average step, std of step
            self.step_mean = 80 # m
            self.step_std = self.step_mean / 4 # m
            self.infectious_radius = self.dlg.box_radius.value()
            self.knn = KNeighborsClassifier(metric='euclidean')

            # Read file situation_initiate
            situation_initiate_file = self.dlg.lineEdit.text()
            situation_initiate = self._read_situation(situation_initiate_file)
            situation_initiate["infectious"]["76026734"] = 10

            self._init_points(selectedLayer, situation_initiate)
            self._put_agent_on_map()

            self.run_steps = number_of_days * 24 * 60 # 1 step = 1 minute
            self.step = 0
            self.timer = QTimer()

            # self._run_simulation()

    def _next_step(self):
        self._move_points(self.baseLayer, self._sus_layer, self._susceptible,
                          step_mean=self.step_mean, step_std=self.step_std)
        self._move_points(self.baseLayer, self._inf_layer, self._infectious,
                          step_mean=self.step_mean, step_std=self.step_std)
        self._move_points(self.baseLayer, self._rec_layer, self._recovered,
                          step_mean=self.step_mean, step_std=self.step_std)

        self._run_infectious()
        self.step += 1

        if self.step >= self.run_steps:
            self.run_steps = 0
            self.timer.stop()

    def _run_infectious(self):
        # list_sus_geo = np.array([x.get_current_geometry_XY() for x in susceptible])
        # list_sus_ids = np.array([x.get_id() for x in susceptible])
        #
        # list_inf_geo = np.array([x.get_current_geometry_XY() for x in infectious])

        list_sus_geo = []
        list_sus_ids = []
        for feat in self._sus_layer.getFeatures():
            geom = feat.geometry()
            list_sus_geo.append([geom.get().x(), geom.get().y()])
            list_sus_ids.append(feat.id())
        list_sus_geo = np.array(list_sus_geo)
        list_sus_ids = np.array(list_sus_ids)

        n_neighbors = self.default_n_neighbor
        if len(list_sus_ids) < n_neighbors:
            n_neighbors = len(list_sus_ids)

        # Get infectious people
        list_inf_geo = []
        list_inf_ids = []
        for feat in self._inf_layer.getFeatures():
            geom = feat.geometry()
            list_inf_geo.append([geom.get().x(), geom.get().y()])
            list_inf_ids.append(feat.id())
        list_inf_geo = np.array(list_inf_geo)
        list_inf_ids = np.array(list_inf_ids)

        self.knn.fit(list_sus_geo, list_sus_ids)
        dist, kneighbors_ids = self.knn.kneighbors(list_inf_geo, n_neighbors=n_neighbors)

        new_infectious = {}
        for inf_id, list_tobe_inf, list_dist in zip(list_inf_ids, kneighbors_ids, dist):
            list_close_contact = list_tobe_inf[list_dist <= self.infectious_radius]
            prob = np.random.choice([0, 1], size=list_close_contact.shape[0],
                                    replace=True, p=[1-self.beta, self.beta])
            list_close_contact = list_close_contact[prob == 1]

            new_infectious[inf_id] = list_sus_ids[list_close_contact].tolist()

        # close_contact_kneighbors_ids = kneighbors_ids[dist <= self.infectious_radius]
        # raise ValueError(str(close_contact_kneighbors_ids))
        # new_infectious = {k:v for k, v in zip(list_inf_ids, close_contact_kneighbors_ids)}
        # close_contact_kneighbors_ids = list_sus_ids[close_contact_kneighbors_ids]
        # # close_contact_kneighbors_ids = np.unique(close_contact_kneighbors_ids).astype(np.int)
        #
        # prob = np.random.choice([0,1], size=close_contact_kneighbors_ids.shape[0],
        #                         replace=True, p=[1-self.beta, self.beta])
        # close_contact_kneighbors_ids = close_contact_kneighbors_ids[prob == 1]

        self._transmit_epidemic(new_infectious)

    def _transmit_epidemic(self, new_infectious):
        self._inf_layer.startEditing()
        self._sus_layer.startEditing()

        delete_ids = []
        for inf_id, list_ids in new_infectious.items():
            for id in list_ids:
                feat = self._sus_layer.getFeature(id)
                agent_id = feat.attribute("id")

                if id in delete_ids:
                    continue

                self._infectious[agent_id] = self._susceptible[agent_id]
                self._infectious[agent_id].set_status(Status.INFECTIOUS)
                self._infectious[agent_id].set_infectious_root(inf_id)
                del self._susceptible[agent_id]

                self._inf_layer.dataProvider().addFeature(feat)
                delete_ids.append(id)

        self._sus_layer.dataProvider().deleteFeatures(delete_ids)

        self._sus_layer.commitChanges()
        self._inf_layer.commitChanges()

    def _run_simulation(self):
        if self.run_steps == 0:
            return

        # simulation = functools.partial(self._next_step,
        #                                baseLayer=self.baseLayer,
        #                                run_steps=self.run_steps,
        #                                step_mean=self.step_mean,
        #                                step_std=self.step_std)
        self.timer.timeout.connect(self._next_step)
        self.timer.start(1500)

    def _construct_adjacent_ward(self):
        index = QgsSpatialIndex()
        for f in self.baseLayer.getFeatures():
            index.addFeature(f)

        for ward in self.baseLayer.getFeatures():
            geom = ward.geometry()
            maphuong = ward.attribute("maphuong")
            self._llist_adjacent_ward[maphuong] = []
            intersecting_ids = index.intersects(geom.boundingBox())

            for intersecting_id in intersecting_ids:
                intersect_ward = self.baseLayer.getFeature(intersecting_id)
                # ward != intersect_ward and
                if (not intersect_ward.geometry().disjoint(geom)):
                    self._llist_adjacent_ward[maphuong].append(intersect_ward.attribute("maphuong"))

    def _get_adjacent_wards(self, maphuong):
        return self._llist_adjacent_ward[maphuong]

    def _read_situation(self, filename):
        situation_initiate = {}

        with open(filename, "r", encoding="utf-8") as f:
            data_read = f.read()
            first_line = True
            id_label_mapping = {}
            for line in data_read.split("\n"):
                if first_line:
                    names = line.split(",")
                    for idx, name in enumerate(names):
                        id_label_mapping[idx] = name
                        if idx > 1:
                            situation_initiate[name] = {}

                    first_line = False
                else:
                    elements = line.split(",")
                    current_ward_id = ""
                    for idx, element in enumerate(elements):
                        current_label = id_label_mapping[idx]

                        if current_label == "maphuong":
                            current_ward_id = element
                        elif idx > 1:
                            situation_initiate[current_label][current_ward_id] = int(element)

        return situation_initiate

    def _move_points(self, baseLayer, layer, agents, step_mean=80, step_std=10):
        layer.startEditing()
        for feat in layer.getFeatures():
            agent_id = feat.attribute("id")
            if agent_id not in agents:
                continue

            agent_current_ward = agents[agent_id].get_current_ward()
            list_consider_ward_ids = self._get_adjacent_wards(agent_current_ward)

            geom = feat.geometry()
            x_coor = geom.get().x()
            y_coor = geom.get().y()

            x_direction = random.gauss(step_mean, step_std) * random.randint(-1, 1)
            y_direction = random.gauss(step_mean, step_std)* random.randint(-1, 1)
            new_position = QgsPointXY(x_coor + x_direction, y_coor + y_direction)

            while True:
                gen_success = False
                for maphuong in list_consider_ward_ids:
                    wid = self._get_map_ward_id(maphuong)
                    ward = baseLayer.getFeature(wid)

                    if ward.geometry().contains(new_position):
                        gen_success = True
                        agent_current_ward = maphuong
                        break

                if gen_success:
                    break
                else:
                    x_direction = random.gauss(step_mean, step_std) * random.randint(-1, 1)
                    y_direction = random.gauss(step_mean, step_std)* random.randint(-1, 1)
                    new_position = QgsPointXY(x_coor + x_direction, y_coor + y_direction)

            geom.translate(x_direction, y_direction)

            agents[agent_id].move(agent_current_ward, new_position)

            layer.changeGeometry(feat.id(), geom)

        layer.commitChanges()

    def _put_agent_on_map(self):
        self._add_point_on_map(self._sus_layer, self._susceptible)
        self._add_point_on_map(self._inf_layer, self._infectious)
        self._add_point_on_map(self._rec_layer, self._recovered)
        self._add_point_on_map(self._dec_layer, self._deceased)

    def _add_point_on_map(self, layer, agents):
        vpr = layer.dataProvider()

        for id, agent in agents.items():
            pnt = QgsGeometry.fromPointXY(agent.get_current_geometry())
            f = QgsFeature()
            f.setGeometry(pnt)
            f.setAttributes([str(id)])

            vpr.addFeature(f)

        layer.updateExtents()
        QgsProject.instance().addMapLayer(layer)

    def _get_map_ward_id(self, maphuong):
        return self._mapping_ward_id[maphuong]

    def _init_points(self, layer, init_situation):
        init_infectious = init_situation["infectious"]
        init_recovered = init_situation["recovered"]
        init_deceased = init_situation["deceased"]

        susceptible = {}
        infectious = {}
        recovered = {}
        deceased = {}

        unique_id = [0]

        for ward in layer.getFeatures():
            ward_geometry = ward.geometry()
            ward_bbox = ward_geometry.boundingBox()
            ward_id = ward.attribute("maphuong")
            # district_id = ward.attribute("maquan")
            ward_population = int(ward.attribute("dan_so").replace(",",""))

            self._mapping_ward_id[ward_id] = ward.id()

            if unique_id[0] > 0:
                continue

            x_min = ward_bbox.xMinimum()
            x_max = ward_bbox.xMaximum()
            y_min = ward_bbox.yMinimum()
            y_max = ward_bbox.yMaximum()
            bbox = (x_min, y_min, x_max, y_max)

            susceptible_people = ward_population - init_infectious[ward_id]
            susceptible.update(self._generate_agent(ward_id,
                                                ward_geometry,
                                                bbox,
                                                90,
                                                Status.SUSCEPTIBLE,
                                                unique_id))

            infectious_people = init_infectious[ward_id]
            infectious.update(self._generate_agent(ward_id,
                                               ward_geometry,
                                               bbox,
                                               infectious_people,
                                               Status.INFECTIOUS,
                                               unique_id))

            recovered_people = init_recovered[ward_id]
            recovered.update(self._generate_agent(ward_id,
                                                ward_geometry,
                                                bbox,
                                                recovered_people,
                                                Status.RECOVERED,
                                                unique_id))

            deceased_people = init_deceased[ward_id]
            deceased.update(self._generate_agent(ward_id,
                                               ward_geometry,
                                               bbox,
                                               deceased_people,
                                               Status.DECEASED,
                                               unique_id))

            # break

        self._susceptible = susceptible
        self._infectious = infectious
        self._recovered = recovered
        self._deceased = deceased

    def _generate_agent(self, ward_id, ward_geometry, bbox, people, status, start_id):
        output = {}
        x_min, y_min, x_max, y_max = bbox

        for _ in range(people):
            agent_x = random.uniform(x_min, x_max)
            agent_y = random.uniform(y_min, y_max)

            while not ward_geometry.contains(QgsPointXY(agent_x, agent_y)):
                agent_x = random.uniform(x_min, x_max)
                agent_y = random.uniform(y_min, y_max)

            output[str(start_id[0])] = Agent(start_id[0],
                                        status,
                                        ward_id,
                                        QgsPointXY(agent_x, agent_y))

            start_id[0] += 1

        return output
