package holeg;

import holeg.model.Grid;
import holeg.model.GridEdge;
import holeg.model.GridNode;
import holeg.model.NodeType;
import holeg.power_flow.*;
import holeg.ui.PowerFlowAnalysisMenu;

import javax.swing.*;
import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.io.PrintStream;
import java.nio.charset.StandardCharsets;
import java.util.*;

public class HolegPowerFlow {
    private Random random = new Random();

    public GridSolverResult solve(Grid grid, PowerFlowSettings settings) {
        if (grid == null)
            throw new IllegalArgumentException("grid is null");
        if (settings == null)
            throw new IllegalArgumentException("settings is null");

        List<Grid> islands = findIslands(grid);

        boolean anySolved = false;
        GridSolverResult sumResult = GridSolverResult.Solved;
        double minSlackPower = Double.MAX_VALUE;

        // Try to solve each island
        for (Grid island : islands) {
            boolean solved = false;

            // Calculate number of calculation tries
            int solveTryCount;
            if (island.getNodes().stream().anyMatch(n -> n.getTypeByDesign() == NodeType.Slack))
                solveTryCount = 1;
            else if (settings.slackNodePlacementStrategy == SlackNodePlacementStrategy.MinimizeSlack)
                solveTryCount = island.getNodes().size();
            else if (settings.slackNodePlacementStrategy == SlackNodePlacementStrategy.RandomNode
                    || settings.slackNodePlacementStrategy == SlackNodePlacementStrategy.RandomProducer)
                solveTryCount = island.getNodes().size() / 2;
            else
                solveTryCount = 1;

            PowerFlowProblem bestProblem = null;
            SolverResult bestResult = null;

            for (int i = 0; i < solveTryCount; i++) {
                // Build problem and solve it
                PowerFlowProblem problem = buildProblem(island, i, settings);

                // See if we can skip the problem
                if (problem.buses.length <= 1) {
                    solved = true;
                    break;
                }
                if (settings.skipGridsWithNoProducers && problem.getProducers().isEmpty()) {
                    solved = true;
                    break;
                }

                Solver solver = new NewtonRaphsonSolver();
                SolverResult result = solver.solve(problem, settings.solverSettings);

                // Debug message
                if (PowerFlowAnalysisMenu.getInstance().shouldShowDebug() && (result.solved || i == solveTryCount - 1))
                    debugShowResultAsMessageBox(problem, result);

                if (result.solved) {
                    solved = true;

                    // Simply use the first solved result
                    if (settings.slackNodePlacementStrategy != SlackNodePlacementStrategy.MinimizeSlack) {
                        bestProblem = problem;
                        bestResult = result;
                        break;
                    }

                    // Is this result better?
                    double slackPower = result.flowData.busInjection[0].lenSquared();
                    if (slackPower < minSlackPower) {
                        minSlackPower = slackPower;
                        bestProblem = problem;
                        bestResult = result;
                    }
                }
            }

            // Update grid if solved
            if (solved) {
                // Check limits
                boolean error = false;
                if (bestResult != null) {
                    if (bestResult.flowData.busInjection[0].len() > settings.maxSlackPowerUntilInvalid)
                        error = true;
                    if (Arrays.stream(bestResult.voltages).anyMatch(u -> u.real < settings.minVoltageUntilInvalid))
                        error = true;
                }

                // Update grid or set error
                if (!error) {
                    if (bestProblem != null)
                        updateGrid(island, bestProblem, bestResult);
                    anySolved = true;
                }
                else
                    sumResult = GridSolverResult.Error;
            }
            else
                sumResult = GridSolverResult.Error;
        }

        // Check if some islands failed, but not all
        if (anySolved && sumResult == GridSolverResult.Error)
            return GridSolverResult.PartialFailure;
        return sumResult;
    }

    private List<Grid> findIslands(Grid grid) {
        List<Grid> islands = new ArrayList<>();
        List<GridNode> nodes = grid.getNodes();

        // Keeps track if the node at any index was already visited
        boolean[] visited = new boolean[nodes.size()];

        for (int i = 0; i < nodes.size(); i++) {
            if (visited[i])
                continue;

            // First node or node that is not connected to any other node -> new island
            Grid island = new Grid();
            Stack<Integer> nodesToVisit = new Stack<>();

            // Add first node
            nodesToVisit.push(i);

            // Depth first search algorithm
            while (!nodesToVisit.empty()) {
                int index = nodesToVisit.pop();

                // Check if already visited
                if (visited[index])
                    continue;
                visited[index] = true;

                // Add node to island
                GridNode node = nodes.get(index);
                island.addNode(node);

                // Iterate over all edges
                List<GridEdge> edges = grid.getEdges(node);
                for (GridEdge edge : edges) {
                    island.addEdge(edge);

                    // Select other end of the edge for next visit
                    if (edge.getFrom() == node)
                        nodesToVisit.push(nodes.indexOf(edge.getTo()));
                    else
                        nodesToVisit.push(nodes.indexOf(edge.getFrom()));
                }
            }
            islands.add(island);
        }
        return islands;
    }

    private PowerFlowProblem buildProblem(Grid grid, int tryCount, PowerFlowSettings settings) {
        List<Bus> buses = new ArrayList<>();
        List<Line> lines = new ArrayList<>();

        // Convert from grid to power flow objects
        for (GridNode node : grid.getNodes())
            createGridNode(buses, node);
        for (GridEdge edge : grid.getEdges())
            createGridLine(grid, lines, edge);

        // Scale voltages and scale power
        double scaleVoltage = 1.0 / 230.0;
        double scalePower = scalePower(buses);

        // Find position for slack node
        addSlackNode(buses, lines, tryCount, settings);

        // Sort lines
        lines.sort(Comparator.comparingInt(a -> a.from));

        // Create problem
        return new PowerFlowProblem(buses.toArray(new Bus[0]), lines.toArray(new Line[0]), scaleVoltage, scalePower);
    }

    private void createGridNode(List<Bus> buses, GridNode node) {
        Bus bus = new Bus();
        bus.delta = 0;
        bus.voltage = 1;
        bus.Pl = Math.abs(node.getPowerConsumption().real);
        bus.Ql = Math.abs(node.getPowerConsumption().imaginary);
        bus.tag = node;

        if (node.getTypeByDesign() == NodeType.Slack)
            bus.type = BusType.Slack;
        else if (node.getPowerGeneration().isZero()) {
            bus.Pg = 0;
            bus.Qg = 0;
            bus.type = BusType.PQ;
        }
        else {
            bus.Pg = Math.abs(node.getPowerGeneration().real);
            bus.Qg = Math.abs(node.getPowerGeneration().imaginary);
            bus.type = BusType.PV;
        }
        bus.Qmin = -1000000;
        bus.Qmax = 1000000;

        buses.add(bus);
    }

    private void createGridLine(Grid grid, List<Line> lines, GridEdge edge) {
        Line line = new Line();
        line.tag = edge;
        if (ComplexNumber.isNullOrZero(edge.getOverrideImpedance())) {
            // Calculate values for line based on sample datasheet
            line.R = 0.000194 * edge.getLengthKilometers();
            line.X = 0.000592 * edge.getLengthKilometers();
            line.B_2 = 0.000528;
        }
        else {
            line.R = edge.getOverrideImpedance().real;
            line.X = edge.getOverrideImpedance().imaginary;
            line.B_2 = edge.getOverrideShuntSusceptance();
        }
        line.a = 1;

        // Calculate from/to, swap if needed
        int from = grid.getNodes().indexOf(edge.getFrom());
        int to = grid.getNodes().indexOf(edge.getTo());
        line.reverse = from > to;
        line.from = Math.min(from, to);
        line.to = Math.max(from, to);

        lines.add(line);
    }

    private double scalePower(List<Bus> buses) {
        double decimal = 1;
        while(true) {
            boolean nextDecimal = false;
            for (Bus bus : buses) {
                if (Math.abs(bus.Pl) > decimal || Math.abs(bus.Pg) > decimal || Math.abs(bus.Ql) > decimal || Math.abs(bus.Qg) > decimal) {
                    nextDecimal = true;
                    break;
                }
            }
            if (!nextDecimal)
                break;
            decimal *= 10;
        }
        for (Bus bus : buses) {
            bus.Pl /= decimal;
            bus.Pg /= decimal;
            bus.Ql /= decimal;
            bus.Qg /= decimal;
            bus.Qmin /= decimal;
            bus.Qmax /= decimal;
        }
        return 1.0 / decimal;
    }

    private void addSlackNode(List<Bus> buses, List<Line> lines, int tryCount, PowerFlowSettings settings) {
        // Skip if have already slack node
        int slackNodePosition = -1;
        for (int i = 0; i < buses.size(); i++) {
            if (buses.get(i).type == BusType.Slack) {
                slackNodePosition = i;
                break;
            }
        }

        // Position
        int position;
        if (slackNodePosition >= 0)
            position = slackNodePosition;
        else if (settings.slackNodePlacementStrategy == SlackNodePlacementStrategy.FirstNode)
            position = 0;
        else if (settings.slackNodePlacementStrategy == SlackNodePlacementStrategy.RandomNode
                || settings.slackNodePlacementStrategy == SlackNodePlacementStrategy.RandomProducer) {

            // Find possible positions
            List<Integer> possiblePositions = new ArrayList<>();
            if (settings.slackNodePlacementStrategy == SlackNodePlacementStrategy.RandomProducer) {
                for (int i = 0; i < buses.size(); i++)
                    if (buses.get(i).isProducer())
                        possiblePositions.add(i);
            }

            // Choose position in grid
            if (possiblePositions.size() == 0) // No generators or strategy not random producer? Add all positions
                position = random.nextInt(buses.size());
            else
                position = possiblePositions.get(random.nextInt(possiblePositions.size()));
        }
        else if (settings.slackNodePlacementStrategy == SlackNodePlacementStrategy.LargestProducer) {
            position = 0;

            // Find producer with max Pg value
            double maxPower = -Double.MAX_VALUE;
            for (int i = 0; i < buses.size(); i++) {
                double power = buses.get(i).Pg;
                if (power > maxPower) {
                    position = i;
                    maxPower = power;
                }
            }
        }
        else if (settings.slackNodePlacementStrategy == SlackNodePlacementStrategy.LargestConsumer) {
            position = 0;

            // Find producer with max Pl value
            double maxPower = -Double.MAX_VALUE;
            for (int i = 0; i < buses.size(); i++) {
                double power = buses.get(i).Pl;
                if (power > maxPower) {
                    position = i;
                    maxPower = power;
                }
            }
        }
        else if (settings.slackNodePlacementStrategy == SlackNodePlacementStrategy.MinimizeSlack) {
            position = tryCount;
            double power = buses.get(position).Pg + buses.get(position).Qg + buses.get(position).Pl + buses.get(position).Ql;
            if (Math.abs(power) < 0.01)
                return;
        }
        else
            throw new UnsupportedOperationException("SlackNodePlacementStrategy not implemented");

        // Already contains slack node or we should replace one node with a new slack node
        if (slackNodePosition >= 0 || settings.replaceNodeWithSlackNode) {
            Bus slack = buses.get(position);
            slack.type = BusType.Slack;

            if (position > 0) {
                // Swap position with first element
                Bus zeroNode = buses.get(0);
                buses.set(0, slack);
                buses.set(position, zeroNode);

                // Fix indices of lines (we changed the order in the list)
                for (Line line : lines) {
                    if (line.from == position)
                        line.from = 0;
                    else if (line.from == 0)
                        line.from = position;

                    if (line.to == position)
                        line.to = 0;
                    else if (line.to == 0)
                        line.to = position;

                    // We need to swap from/to
                    if (line.from > line.to) {
                        int t = line.from;
                        line.from = line.to;
                        line.to = t;
                        line.reverse = !line.reverse;
                    }
                }
            }
        }
        else {
            // Create slack node
            Bus slack = new Bus();
            slack.voltage = 1;
            slack.delta = 0;
            slack.type = BusType.Slack;

            buses.add(0, slack);

            // Fix line indices
            for (Line line : lines) {
                line.from += 1;
                line.to += 1;
            }

            // Create short line
            Line line = new Line();
            line.from = position + 1;
            line.to = buses.indexOf(slack);
            line.R = 0.01938;
            line.X = 0.1;
            line.B_2 = 0.1;
            line.a = 1;

            lines.add(line);
        }
    }

    private void debugShowResultAsMessageBox(PowerFlowProblem problem, SolverResult result) {
        try (ByteArrayOutputStream stream = new ByteArrayOutputStream()) {
            final String utf8 = StandardCharsets.UTF_8.name();
            try (PrintStream ps = new PrintStream(stream, true, utf8)) {
                result.printTo(problem, ps);
            }
            String data = stream.toString(utf8);
            JOptionPane.showMessageDialog(null, data,"Debug output", JOptionPane.INFORMATION_MESSAGE);
        } catch (IOException e) {
            e.printStackTrace();
        }
    }

    private void updateGrid(Grid grid, PowerFlowProblem problem, SolverResult result) {
        double currentScale = 1;
        for (int i = 0; i < problem.buses.length; i++) {
            GridNode node = (GridNode)problem.buses[i].tag;
            if (node == null)
                continue;

            node.setVoltage(result.voltages[i].len() / problem.scaleVoltage);
            node.setPhase(Math.toDegrees(result.voltages[i].angle()));
            node.setCurrent(result.flowData.busCurrent[i] / currentScale);

            // Overwrite slack type
            if (problem.buses[i].type == BusType.Slack)
                node.setTypeSolved(NodeType.Slack);
            else
                node.setTypeSolved(node.getTypeByDesign());
        }

        for (int i = 0; i < problem.lines.length; i++) {
            GridEdge edge = (GridEdge)problem.lines[i].tag;
            if (edge == null)
                continue;

            double reverseFactor = problem.lines[i].reverse ? -1 : 1;
            edge.setPowerFlow(result.flowData.linePower[i].multiply(1 / problem.scalePower).multiply(reverseFactor));
            edge.setLineLoss(result.flowData.linePowerLoss[i].multiply(1 / problem.scalePower).multiply(reverseFactor));
            edge.setCurrent(result.flowData.lineCurrent[i] / currentScale * reverseFactor);
        }
    }
}