#  Copyright 2020 Zeppelin Bend Pty Ltd
#
#  This Source Code Form is subject to the terms of the Mozilla Public
#  License, v. 2.0. If a copy of the MPL was not distributed with this
#  file, You can obtain one at https://mozilla.org/MPL/2.0/.

from __future__ import annotations

import copy
import logging
from dataclassy import dataclass
from enum import Enum


from zepben.evolve.model import EnergySource
from zepben.evolve.model import Breaker
from zepben.evolve.model import SinglePhaseKind
from zepben.evolve.model.phasedirection import PhaseDirection
from zepben.evolve.exceptions import PhaseException
from zepben.evolve.services.network.tracing.connectivity import get_connectivity
from zepben.evolve.exceptions import TracingException
from zepben.evolve.services.network.tracing.phase_status import normal_phases, current_phases
from zepben.evolve.services.network.tracing import queue_next_terminal
from zepben.evolve.services.network.tracing.traversals.queue import PriorityQueue
from zepben.evolve.services.network.tracing.traversals.tracing import Traversal
from zepben.evolve.services.network.tracing.phase_status import PhaseStatus
from zepben.evolve.services.network.tracing.traversals import BranchRecursiveTraversal
from zepben.evolve.services.network.tracing.util import normally_open, currently_open
from typing import Set, Callable, List, Iterable, Optional

__all__ = ["FeederProcessingStatus", "SetPhases", "FeederCbTerminalPhasesByStatus", "DelayedFeederTrace",
           "set_phases_and_queue_next", "set_current_phases_and_queue_next", "set_normal_phases_and_queue_next"]

logger = logging.getLogger("phasing.py")


class FeederProcessingStatus(Enum):
    COMPLETE = 0,
    PARTIAL = 1,
    NONE = 2


@dataclass(slots=True)
class FeederCbTerminalPhasesByStatus:
    terminal: Terminal
    in_phases: Set[SinglePhaseKind] = set()
    none_phases: Set[SinglePhaseKind] = set()
    phases_to_flow: Set[SinglePhaseKind] = set()


@dataclass(slots=True)
class DelayedFeederTrace:
    out_terminal: Terminal
    phases_to_flow: Set[SinglePhaseKind]


class SetPhases(object):
    def __init__(self):
        self.normal_traversal = BranchRecursiveTraversal(queue_next=set_normal_phases_and_queue_next,
                                                         process_queue=PriorityQueue(),
                                                         branch_queue=PriorityQueue())
        self.current_traversal = BranchRecursiveTraversal(queue_next=set_current_phases_and_queue_next,
                                                          process_queue=PriorityQueue(),
                                                          branch_queue=PriorityQueue())

    async def run(self, network: NetworkService):
        # terminals = await _apply_phases_from_feeder_cbs(network)
        await _apply_phases_from_sources(network)
        terminals = [term for es in network.objects(EnergySource) if es.num_phases() > 0 for term in es.terminals]
        if not terminals:
            raise TracingException("No feeder sources were found, tracing cannot be performed.")
        breakers = network.objects(Breaker)
        await self.run_complete(terminals, breakers)

    async def run_complete(self, terminals: Iterable[Terminal], breakers: Iterable[Breaker]):
        feeder_cbs = [br for br in breakers if br.is_substation_breaker()]
        await self._run_normal(terminals, feeder_cbs)
        await self._run_current(terminals, feeder_cbs)

    async def _run_normal(self, terminals, feeder_cbs):
        await run_set_phasing(terminals, feeder_cbs, self.normal_traversal, normally_open, normal_phases)

    async def _run_current(self, terminals, feeder_cbs):
        await run_set_phasing(terminals, feeder_cbs, self.current_traversal, currently_open, current_phases)

    async def run_ce(self, ce: ConductingEquipment, breakers: Iterable[Breaker]):
        if ce.num_terminals() == 0:
            return
        for in_term in ce.terminals:
            normal_phases_to_flow = _get_phases_to_flow(in_term, normally_open, normal_phases)
            current_phases_to_flow = _get_phases_to_flow(in_term, currently_open, current_phases)

            for out_term in ce.terminals:
                if out_term is not in_term:
                    _flow_through_equipment(self.normal_traversal, in_term, out_term, normal_phases_to_flow, normal_phases)
                    _flow_through_equipment(self.current_traversal, in_term, out_term, current_phases_to_flow, current_phases)

        self.normal_traversal.tracker.clear()
        self.current_traversal.tracker.clear()

        await self.run_complete(ce.terminals, breakers)


async def find_es_breaker_terminal(es):
    """
    From an EnergySource finds the closest connected Feeder CB (Breaker that is part of a substation).
    At the moment we assume that all EnergySource's with EnergySourcePhase's will be associated with at least a
    single feeder circuit breaker, and thus this function given an `EnergySource` will perform a trace that returns
    the first `zepben.evolve.iec61970.base.core.terminal.Terminal` encountered from that `EnergySource` that belongs to a `Breaker`. This `zepben.evolve.iec61970.base.core.terminal.Terminal` should always
    be the most downstream `zepben.evolve.iec61970.base.core.terminal.Terminal` on the `Breaker`, and thus can then be used for setting `Direction` downstream and
    away from this `Breaker`.
    TODO: check how ES are normally connected to feeder CB's.
    """
    out_terminals = set()

    async def stop_on_sub_breaker(term, exc=None):
        if out_terminals:  # stop as soon as we find a substation breaker.
            return True
        try:
            if term.conducting_equipment.is_substation_breaker():
                out_terminals.add(term)
                return True
        except AttributeError:
            return False
        return False

    t = Traversal(queue_next=queue_next_terminal, start_item=es.terminals[0], process_queue=PriorityQueue(), stop_conditions=[stop_on_sub_breaker])
    await t.trace()

    return out_terminals


async def _apply_phases_from_feeder_cbs(network):
    """
    Apply phase and direction on all Feeder Circuit Breakers. Will make all phases on the outgoing Terminal of a
    `Breaker` that is part of a substation have a `Direction` of `OUT`.
    `network` `zepben.evolve.network.Network` to apply phasing on.
    """
    start_terms = []
    # TODO: check if below assumption is correct
    # We find the substation breaker from the networks energy sources as we assume that the ES will be wired below
    # the breaker, and thus we can determine which terminal of the breaker to flow out from and apply phases.
    for es in network.energy_sources.values():
        esp = es.energy_source_phases
        if esp:
            if len(esp) != es.num_cores:
                # TODO: java network phases doesn't throw here, but would throw in the below for loop if num_phases > len(esp). why does java silently handle less cores?
                logger.error(f"Energy source {es.name} [{es.mrid}] is a source with {len(esp)} and {es.num_phases}. Number of phases should match number of cores. Phasing cannot be applied")
                raise TracingException(
                    f"Energy source {es.name} [{es.mrid}] is a source with {len(esp)} and {es.num_cores}. Number of phases should match number of cores. Phasing cannot be applied")
            breaker_terms = await find_es_breaker_terminal(es)
            for terminal in breaker_terms:
                for phase in terminal.phases.single_phases:
                    normal_phases(terminal, phase).add(esp[phase].phase, PhaseDirection.OUT)
                    current_phases(terminal, phase).add(esp[phase].phase, PhaseDirection.OUT)
                logger.debug(f"Set {terminal.conducting_equipment.mrid} as Feeder Circuit Breaker with phases {terminal.phases.phase}")
            start_terms.extend(breaker_terms)
    return start_terms


async def _apply_phases_from_sources(network: NetworkService):
    """
    Apply phase and direction on all Feeder Circuit Breakers. Will make all phases on the outgoing Terminal of a
    `Breaker` that is part of a substation have a `Direction` of `OUT`.
    `network` `zepben.evolve.network.Network` to apply phasing on.
    """
    for es in network.objects(EnergySource):
        if es.num_phases() > 0:
            await _apply_phases_from_source(es)


async def _apply_phases_from_source(energy_source: EnergySource):
    if energy_source.num_terminals() == 0:
        return
    es_phases = set()
    for phase in energy_source.phases:
        es_phases.add(phase)

    nominal_phases = set()
    for terminal in energy_source.terminals:
        nominal_phases.update(terminal.phases.single_phases)

    if len(es_phases) != len(nominal_phases):
        logger.warning((f"Energy source {str(energy_source)} is a source with {len(es_phases)} phases and {len(nominal_phases)} nominal phases. "
                        f"Number of phases should match the number of nominal phases!"))

    for term in energy_source.terminals:
        for phase in term.phases.single_phases:
            normal_phases(term, phase).add(phase, PhaseDirection.OUT)
            current_phases(term, phase).add(phase, PhaseDirection.OUT)


# TODO: pass through visited and be smart with it
def set_normal_phases_and_queue_next(terminal, traversal, visited):
    set_phases_and_queue_next(terminal, traversal, normally_open, normal_phases)


def set_current_phases_and_queue_next(terminal, traversal, visited):
    set_phases_and_queue_next(terminal, traversal, currently_open, current_phases)


def set_phases_and_queue_next(current: Terminal,
                              traversal: BranchRecursiveTraversal,
                              open_test: Callable[[ConductingEquipment, SinglePhaseKind], bool],
                              phase_selector: Callable[[Terminal, SinglePhaseKind], PhaseStatus]):
    phases_to_flow = _get_phases_to_flow(current, open_test, phase_selector)

    if current.conducting_equipment:
        for out_terminal in current.conducting_equipment.terminals:
            if out_terminal != current and _flow_through_equipment(traversal, current, out_terminal, phases_to_flow, phase_selector):
                _flow_out_to_connected_terminals_and_queue(traversal, out_terminal, phases_to_flow, phase_selector)


async def run_set_phasing(start_terminals: List[Terminal],
                          process_feeder_cbs: List[Breaker],
                          traversal: BranchRecursiveTraversal,
                          open_test: Callable[[ConductingEquipment, Optional[SinglePhaseKind]], bool],
                          phase_selector: Callable[[Terminal, Optional[SinglePhaseKind]], PhaseStatus]):
    for terminal in start_terminals:
        await _run_terminal(terminal, traversal, phase_selector)

    # We take a copy of the feeder CB's as we will modify the list while processing them.
    process_feeder_cbs = copy.copy(process_feeder_cbs)
    keep_processing = True
    while keep_processing:
        delayed_feeder_traces = []
        for feeder_cb in process_feeder_cbs:
            status = _run_feeder_breaker(feeder_cb, traversal, open_test, phase_selector, delayed_feeder_traces)
            if status == FeederProcessingStatus.COMPLETE:
                process_feeder_cbs.remove(feeder_cb)

        for trace in delayed_feeder_traces:
            await _run_from_out_terminal(traversal, trace.out_terminal, trace.phases_to_flow, phase_selector)
        keep_processing = len(delayed_feeder_traces) > 0


async def _run_terminal(start: Terminal, traversal: BranchRecursiveTraversal, phase_selector: Callable[[Terminal, SinglePhaseKind], PhaseStatus]):
    phases_to_flow = {phase for phase in start.phases.single_phases if phase_selector(start, phase).direction().has(PhaseDirection.OUT)}
    await _run_from_out_terminal(traversal, start, phases_to_flow, phase_selector)


async def _run_from_out_terminal(traversal: BranchRecursiveTraversal, out_terminal: Terminal, phases_to_flow: Set[SinglePhaseKind],
                                 phase_selector: Callable[[Terminal, SinglePhaseKind], PhaseStatus]):
    traversal.reset()
    traversal.tracker.visit(out_terminal)
    _flow_out_to_connected_terminals_and_queue(traversal, out_terminal, phases_to_flow, phase_selector)
    await traversal.trace()


def _flow_out_to_connected_terminals_and_queue(traversal: BranchRecursiveTraversal, out_terminal: Terminal, phases_to_flow: Set[SinglePhaseKind],
                                               phase_selector: Callable[[Terminal, SinglePhaseKind], PhaseStatus]):
    connectivity_results = get_connectivity(out_terminal, phases_to_flow)
    for cr in connectivity_results:
        in_term = cr.to_terminal
        has_added = False
        for oi in cr.nominal_phase_paths:
            out_core = oi.from_core
            in_core = oi.to_core
            out_phase = phase_selector(out_terminal, out_core).phase()
            in_phase = phase_selector(in_term, in_core)
            try:
                if in_phase.add(out_phase, PhaseDirection.IN):
                    has_added = True
            except PhaseException as ex:
                raise PhaseException(
                    (f"Attempted to apply more than one phase to [{in_term.conducting_equipment.mrid if in_term.conducting_equipment else in_term.mrid}"
                     f" on nominal phase {oi.to_phase}. Attempted to apply phase {out_phase} to {in_phase.phase()}."), ex)

        if has_added and not traversal.has_visited(in_term):
            if len(connectivity_results) > 1 or (out_terminal.conducting_equipment is not None and out_terminal.conducting_equipment.num_terminals() > 2):
                branch = traversal.create_branch()
                branch.start_item = in_term
                traversal.branch_queue.put(branch)
            else:
                traversal.process_queue.put(in_term)


def _get_phases_to_flow(terminal: Terminal,
                        open_test: Callable[[ConductingEquipment, Optional[SinglePhaseKind]], bool],
                        phase_selector: Callable[[Terminal, SinglePhaseKind], PhaseStatus]):
    phases_to_flow = set()
    try:
        if terminal.conducting_equipment.is_substation_breaker():
            return phases_to_flow
    except AttributeError:
        pass

    if terminal.conducting_equipment is None:
        return phases_to_flow

    equip = terminal.conducting_equipment
    for phase in terminal.phases.single_phases:
        if not open_test(equip, phase) and phase_selector(terminal, phase).direction().has(PhaseDirection.IN):
            phases_to_flow.add(phase)
    return phases_to_flow


def _flow_through_equipment(traversal: BranchRecursiveTraversal, in_terminal: Terminal, out_terminal: Terminal, phases_to_flow: Set[SinglePhaseKind],
                            phase_selector: Callable[[Terminal, SinglePhaseKind], PhaseStatus]):
    traversal.tracker.visit(out_terminal)
    has_changes = False

    for phase in phases_to_flow:
        out_phase_status = phase_selector(out_terminal, phase)
        try:
            in_phase = phase_selector(in_terminal, phase).phase()
            applied = out_phase_status.add(in_phase, PhaseDirection.OUT)
            has_changes = applied or has_changes
        except PhaseException as ex:
            raise PhaseException((
                f"Attempted to apply more than one phase to {out_terminal.conducting_equipment.mrid if out_terminal.conducting_equipment else in_terminal.mrid}"
                f" on nominal phase {phase}. Detected phases {out_phase_status.phase()} and {in_phase}. Underlying error was {str(ex)}"),
                ex)
    return has_changes


def _get_feeder_cb_terminal_cores_by_status(feeder_cb: Breaker,
                                            open_test: Callable[[ConductingEquipment, Optional[SinglePhaseKind]], bool],
                                            phase_selector: Callable[[Terminal, SinglePhaseKind], PhaseStatus]):
    res = []
    for terminal in feeder_cb.terminals:
        status = FeederCbTerminalPhasesByStatus(terminal=terminal)
        res.append(status)

        for phase in terminal.phases.single_phases:
            phase_status = phase_selector(terminal, phase)
            if phase_status.direction() == PhaseDirection.IN:
                status.in_phases.add(phase)
                if not open_test(feeder_cb, phase):
                    status.phases_to_flow.add(phase)
            elif phase_status.direction() == PhaseDirection.BOTH:
                status.in_phases.add(phase)
            elif phase_status.direction() == PhaseDirection.NONE:
                status.none_phases.add(phase)
    return res


def _flow_through_feeder_cb_and_queue(in_terminal: FeederCbTerminalPhasesByStatus,
                                      out_terminal: FeederCbTerminalPhasesByStatus,
                                      traversal: BranchRecursiveTraversal,
                                      phase_selector: Callable[[Terminal, SinglePhaseKind], PhaseStatus],
                                      delayed_traces: List,
                                      processed_phases: Set[SinglePhaseKind]):
    if not in_terminal.in_phases:
        return

    phases_to_flow = copy.copy(in_terminal.phases_to_flow)
    for phase in in_terminal.terminal.phases.single_phases:
        if phase in in_terminal.in_phases:
            processed_phases.add(phase)

            # Remove any phases that have already been processed from the other side
            if phase not in out_terminal.none_phases:
                phases_to_flow.remove(phase)

    if _flow_through_equipment(traversal, in_terminal.terminal, out_terminal.terminal, phases_to_flow, phase_selector):
        delayed_traces.append(DelayedFeederTrace(out_terminal.terminal, phases_to_flow))


def _run_feeder_breaker(feeder_cb: Breaker,
                        traversal: BranchRecursiveTraversal,
                        open_test: Callable[[ConductingEquipment, Optional[SinglePhaseKind]], bool],
                        phase_selector: Callable[[Terminal, SinglePhaseKind], PhaseStatus],
                        delayed_traces: List):
    if feeder_cb.num_terminals() not in (1, 2):
        logger.warning(f"Ignoring feeder CB {str(feeder_cb)} with {feeder_cb.num_terminals()} terminals, expected 1 or 2 terminals")
        return FeederProcessingStatus.COMPLETE

    if feeder_cb.num_terminals() == 1:
        set_phases_and_queue_next(next(feeder_cb.terminals), traversal, open_test, phase_selector)
        return FeederProcessingStatus.COMPLETE

    processed_phases = set()
    statuses = _get_feeder_cb_terminal_cores_by_status(feeder_cb, open_test, phase_selector)
    _flow_through_feeder_cb_and_queue(statuses[0], statuses[1], traversal, phase_selector, delayed_traces, processed_phases)
    _flow_through_feeder_cb_and_queue(statuses[1], statuses[0], traversal, phase_selector, delayed_traces, processed_phases)

    nominal_phases = {phase for term in feeder_cb.terminals for phase in term.phases.single_phases}
    if len(processed_phases) == len(nominal_phases):
        return FeederProcessingStatus.COMPLETE
    elif processed_phases:
        return FeederProcessingStatus.PARTIAL
    else:
        return FeederProcessingStatus.NONE
