# -*- coding: utf-8

import datetime
import logging
import multiprocessing as mp
import os

import numpy as np
import pandas as pd
from bokeh.layouts import column
from bokeh.models import DatetimeTickFormatter, Div
from bokeh.palettes import Spectral10
from bokeh.plotting import gridplot, figure, output_file, save
from tespy.tools import logger
from tespy.tools.global_vars import err

from digripy.digripynetwork import DigripyNetwork
from digripy.tools import draw_graph

# Only require cupy if Cuda shall be used
try:
    import cupy as cu
except ModuleNotFoundError:
    cu = None


class Simulation:
    def __init__(self, pipe_data, sim_settings, file_prefix, first_frame,
                 display_graph, run_name, output_dir, input_dir,
                 demands, min_demand: float, log_file: str):
        self.in_dir = input_dir
        self.out_dir = output_dir
        self.run_name = run_name
        self.sim_settings = sim_settings
        self.pipe_data = pipe_data

        if cu is None or 'no_cuda' in self.sim_settings and \
                self.sim_settings['no_cuda'] is True:
            self.use_cuda = False
        else:
            self.use_cuda = True

        # Create the actual district heating network
        self.dh_network = self.create_network()

        self._time = None
        self.demands = demands
        self.demands[self.demands < min_demand] = min_demand
        self.demands_mean_max = self.demands.mean(axis=1).max(axis=0)
        self.first_frame = first_frame
        self.total_frames = len(self.demands.index)
        self.res_file_name = f'{file_prefix}_{self.run_name}'
        self.res_file_path = os.path.join(self.out_dir, self.res_file_name)
        logger.define_logging(log_path=True, log_version=True,
                              logpath=self.out_dir, logfile=log_file,
                              screen_level=logging.INFO,
                              file_level=logging.WARNING,
                              file_format=f'Sub-process {self.run_name} - '
                                          f'%(asctime)s-%(levelname)s-'
                                          f'%(message)s)',
                              screen_format=f'Sub-process {self.run_name} - '
                                          f'%(asctime)s-%(levelname)s-'
                                          f'%(message)s')  # - '
                                          # f'%(filename)s %(lineno)d')

        # Define Column titles for plotting of overall system results and
        # create the corresponding data structure
        self.sys_res_cols = {'Heat [W]': ['Demand Heat Source',
                                          'Total Consumer Heat Demand',
                                          'Total Heat Losses'],
                             'Temperature [°C]': ['Heat Source input',
                                                  'Heat Source output'],
                             'Pump Power [W]': ['Demand Pump Power'],
                             'Pressure [bar]]': ['Pump input', 'Pump output',
                                                 'Pressure Drop'],
                             'Velocity [m/s]': ['Heat Source'],
                             'Relative Data [%]': ['Relative Heat Losses '
                                                   '(percent of demands)']}
        cols = [y for x in self.sys_res_cols.values() for y in x]
        self.sys_res = pd.DataFrame(index=self.demands.index, columns=cols,
                                    dtype=float)

        # Define Column titles for plotting of consumer results and create the
        # corresponding data structure
        cols = ['Heat demand [W]', 'Velocity [m/s]', 'Supply Temperature [°C]',
                'Control Valve Aperture [%]']
        idx = pd.MultiIndex.from_product([self.demands.columns,
                                          self.demands.index])
        self.cons_res = pd.DataFrame(index=idx, columns=cols, dtype=float)

        if display_graph:
            draw_graph(self.dh_network.nw)

    def create_network(self):
        return DigripyNetwork(
            pipe_data=self.pipe_data, sim_settings=self.sim_settings,
            in_dir=self.in_dir)

    def run_process(self):
        """
        This function starts the simulation process and triggers the method
        running the frame simulations.
        :return: Returns an array of size 2 that hast been generated by
                 generate_output_files() and contains information and a
                 Pandas Dataframe with the simulation results.
        """
        run_start = datetime.datetime.now()
        process_result = self.sim_settings
        logging.debug(self.sim_settings)
        process_result['result'] = f'UNHANDLED error occured!'
        n = 0
        try:
            if self.use_cuda:
                if cu.is_available():
                    logging.info(f'Using CUDA device '
                                 f'{cu.cuda.get_device_id()}.')
                else:
                    msg = "CUDA shall be used, but is NOT available."
                    logging.error(msg)
                    raise mp.ProcessError(msg)
            for n in range(self.total_frames):
                self.run_sub_frame(n, 1, 1)
            process_result = self.generate_output_files()
        except mp.ProcessError as e:
            logging.error(e)
            process_result = self.generate_output_files()
            process_result['result'] = f'Error in frame {n + self.first_frame}'
        finally:
            run_end = datetime.datetime.now()
            runtime = run_end - run_start
            logging.info(f'Simulation runtime: {runtime}')
            return [process_result, self.sys_res]

    def run_sub_frame(self, frame_no: int, sub_frame_no: int = 1,
                      subs_per_frame: int = 1, is_retry=False):
        """
        Runs a time step or a sub-step for reducing the step size for the
        solver. If a simulation returned a ValueError the step is split into
        several smaller steps until a maximum number of sub-steps has been
        reached. Results are only gathered when running the last sub-step and
        if the frame has not been run successfully before. Method is
        called recursively.
        :param frame_no: Time series frame number to run.
        :param sub_frame_no: Sub-frame number
        :param subs_per_frame: Amount of sub-frames the frame shall be split
        into.
        :param is_retry: Has this frame already been run successfully?
        """

        # The amount of sub-frames that will be added to the current number if
        # a sub-frames fails to run
        increase_sub_frames_by = 3

        # Maximum number of sub-frames the current frame wil be split into
        max_sub_frames = 20

        sub_frame_name = f'{self.first_frame + frame_no} - ' \
                         f'{self.demands.index[frame_no]} - sub-frame ' \
                         f'{sub_frame_no} / {subs_per_frame}'

        logging.info(f'Running frame {sub_frame_name}...')

        # If current run is a retry: don't use previous run's data for
        # initialization, as this tends to be corrupted. Rebuild the network
        # instead
        if is_retry:
            self.dh_network = self.create_network()

        try:
            self.set_cons_demands(frame_no, sub_frame_no,
                                  subs_per_frame)
            if is_retry:
                self.dh_network.nw.solve(mode='design', init_previous=False,
                                         use_cuda=self.use_cuda,
                                         always_all_equations=True)
            else:
                self.dh_network.nw.solve(mode='design', use_cuda=self.use_cuda,
                                         always_all_equations=True)

            # Throw a ValueError if tespy solver did not converge
            if not self.dh_network.nw.progress or self.dh_network.nw.iter == \
                    self.dh_network.nw.max_iter - 1:
                raise ValueError('Calculation does not converge.')

            # Gather results if not in an intermediate step and the frame has
            # not been run before. Else run next
            # sub-frame
            if sub_frame_no == subs_per_frame:
                if not is_retry:
                    self.gather_frame_results()
            else:
                logging.debug(f'Running sub-frame {sub_frame_no + 1}/'
                              f'{subs_per_frame}...')
                self.run_sub_frame(frame_no, sub_frame_no + 1,
                                   subs_per_frame)

        # Deal with simulation runs that resulted in invalid values at some
        # place
        except ValueError or ZeroDivisionError as e:
            logging.warning(f'Error running time step {sub_frame_name}: {e}')

            # If the frame has not been split into as many sub-frames as are
            # allowed the number of sub-frames will be increased and a new run
            # will start beginning with sub-frame 1. Before this the last frame
            # is run again to
            # provide valid starting values for the sub-frames.
            if subs_per_frame < max_sub_frames:
                if frame_no > 0:
                    logging.warning(f'Running last frame again again...')
                    self.run_sub_frame(frame_no - 1, 1, 1, is_retry=True)
                logging.warning(f'Trying '
                                f'{subs_per_frame + increase_sub_frames_by} '
                                f'sub-frames after failing with '
                                f'{sub_frame_name}')
                self.run_sub_frame(frame_no, 1, subs_per_frame +
                                   increase_sub_frames_by)
            else:
                msg = f'Simulation Process has been terminated as it did not' \
                      f' converge in frame {sub_frame_name}.'
                logging.error(msg)
                raise mp.ProcessError(msg)
        except Exception as e:
            logging.error('UNHANDLED ERROR: ', e)

    def set_cons_demands(self, frame_number, sub_frame, subs_per_frame):
        # Heat demands
        frame_demands = self.demands.iloc[frame_number]
        self._time = frame_demands.name
        if subs_per_frame > 1:
            if frame_number > 0:
                previous_demands = self.demands.iloc[frame_number - 1]
            else:
                previous_demands = np.zeros_like(frame_demands)
            sub_frame_demands = frame_demands - (frame_demands -
                                                 previous_demands) * \
                                (1 - sub_frame / subs_per_frame)
        else:
            sub_frame_demands = frame_demands

        logging.debug(f'Consumer Demands: {sub_frame_demands}')

        # Interpolate consumers' demanded temperature, source temperature and
        # the pump pressure following the curves defined settings CSV-file

        # Heat source out and consumer out temperatures
        current_demand_ratio = sub_frame_demands.mean() / self.demands_mean_max
        if current_demand_ratio > self.sim_settings['end_high_t'] / 100:
            t_source = self.sim_settings['t_source_upper_limit']
            t_cons_out = self.sim_settings['t_consumer_return_high']
        elif current_demand_ratio < self.sim_settings['begin_low_t'] / 100:
            t_source = self.sim_settings['t_source_lower_limit']
            t_cons_out = self.sim_settings['t_consumer_return_low']
        else:
            co_source = \
                np.polyfit([self.sim_settings['begin_low_t'] / 100,
                            self.sim_settings['end_high_t'] / 100],
                           [self.sim_settings['t_source_lower_limit'],
                            self.sim_settings['t_source_upper_limit']], 1)
            co_cons = \
                np.polyfit([self.sim_settings['begin_low_t'] / 100,
                            self.sim_settings['end_high_t'] / 100],
                           [self.sim_settings['t_consumer_return_low'],
                            self.sim_settings['t_consumer_return_high']], 1)
            t_source = np.polyval(co_source, current_demand_ratio)
            t_cons_out = np.polyval(co_cons, current_demand_ratio)

        logging.debug(f'Current mean demand is '
                      f'{current_demand_ratio * 100}% of maximum mean.')

        # Pump pressure
        self.dh_network.set_pump_p_out(
            self.sim_settings['p_pump_low'] + current_demand_ratio *
            (self.sim_settings['p_pump_high'] -
             self.sim_settings['p_pump_low']))
        self.dh_network.set_src_t_out(t_source)

        self.dh_network.set_cons_demands(
            demands=sub_frame_demands, t_out=t_cons_out)

    def generate_output_files(self):
        """
        Generates two output files: one xlsx file with the process results and
        an html, that shows some of that data using the Bokeh framework.
        :return: Returns an array of size 2 that contains information and a
                 Pandas Dataframe with the simulation results
        """
        plots = []
        x = self.sys_res.index
        output_data_frame = pd.DataFrame(index=x)
        i = 0
        line_width = 2
        for title, head in self.sys_res_cols.items():
            plot = figure(width=800, height=300, x_axis_type='datetime',
                          title=title)
            for j, name in enumerate(head):
                y = self.sys_res.loc[:, name].values
                plot.line(x=x, y=y, legend_label=name,
                          line_color=Spectral10[j], line_width=line_width,
                          line_alpha=0.9)
                output_data_frame[name] = y
                i += 1
            plot.xaxis.formatter = DatetimeTickFormatter(days=['%B %d'],
                                                         hours=['%H:%M'])
            plot.legend.location = "top_right"
            plot.legend.click_policy = "hide"
            plots.append(plot)

        for title in self.cons_res.columns:
            plot = figure(width=1800, height=300, x_axis_type='datetime',
                          title=title)
            for i, c in enumerate(self.cons_res.index.levels[0]):
                y = self.cons_res.loc[c, :][title]
                cons = c
                output_data_frame[title + ' ' + cons] = y
                plot.line(x=x, y=y, legend_label=cons,
                          line_color=Spectral10[i], line_width=line_width,
                          line_alpha=0.9)
            plot.xaxis.formatter = DatetimeTickFormatter(days=['%B %d'],
                                                         hours=['%H:%M'])
            plot.legend.location = "top_right"
            plot.legend.click_policy = "hide"
            plots.append(plot)

        output_file(self.res_file_path + ".html")
        str_settings = ""
        for key, value in self.sim_settings.items():
            str_settings = str_settings + key + ' : ' + str(value) + " <br>"

        rel_heat_losses = self.sys_res.loc[:, 'Total Heat Losses'].sum() / \
                      self.sys_res.loc[:, 'Total Consumer Heat Demand'].sum()
        rel_pump_power = self.sys_res.loc[:, 'Demand Pump Power'].sum() / \
                     self.sys_res.loc[:, 'Total Consumer Heat Demand'].sum()

        res_out = self.sim_settings
        res_out['id'] = self.res_file_name
        res_out['relative_heat_losses'] = rel_heat_losses
        res_out['relative_pump_power'] = rel_pump_power
        res_out['tolerable_error'] = err
        res_out['result'] = 'success'

        save(column(Div(text=f'{datetime.datetime.now().strftime("%H-%M-%S")} '
                             f'<br> First Frame: {self.demands.index[0]}  '
                             f'<br> Last Frame: {self.demands.index[-1]}  '
                             f'<br> Simulation parameters: <br>{str_settings} '
                             f'<br> Relative heat losses: '
                             f'{round(rel_heat_losses * 100, 1)}% <br> '
                             f'Relative pump power: '
                             f'{round(rel_pump_power * 100, 1)}% <br>'),
                    gridplot(children=plots, sizing_mode='scale_width',
                             ncols=1)))
        with pd.ExcelWriter(self.res_file_path + ".xlsx") as xls_writer:
            output_data_frame.to_excel(xls_writer)
        return res_out

    def gather_frame_results(self):
        """
        Triggered after each successfully run time step, this method gathers
        the results and saves them in self.cons_res or self.sys_res
        """
        total_demands = 0
        consumers = self.dh_network.get_cons_df()

        for c in consumers.itertuples():
            d = c.object.Q.val
            total_demands += d
            self.cons_res.loc[str(c.Index), self._time] = \
                [-d, c.out_con.v.val / ((c.object.D.val / 2)**2) / np.pi,
                 c.in_con.T.val, c.valve.pr.val * 100]

        losses = self.dh_network.get_total_pipe_heat_losses()

        source_demands = self.dh_network.get_src_heat_demand()

        self.sys_res.loc[self._time] = [source_demands,
                                        -total_demands,
                                        -losses,
                                        self.dh_network.get_src_t_in(),
                                        self.dh_network.get_src_t_out(),
                                        self.dh_network.get_pump_power(),
                                        self.dh_network.get_pump_p_in(),
                                        self.dh_network.get_pump_p_out(),
                                        self.dh_network.get_p_drop(),
                                        self.dh_network.get_src_flow_speed(),
                                        # To avoid divide by zero relative
                                        # losses are set to 0 if demands
                                        # are 0
                                        0 if total_demands > -10
                                        else 100 * np.min(
                                            [losses / source_demands, 1])]

        excess_percentage = \
            (total_demands / source_demands * -1 - 1) * 100
        if excess_percentage > 1:
            msg = f'Sum of consumer ({int(-total_demands)}W is ' \
                  f'greater than heat source demand ' \
                  f'({int(source_demands)}W)! ({excess_percentage}%)'
            logging.error(msg)
            raise ValueError(msg)
