package holeg;

import holeg.model.Grid;
import holeg.model.GridEdge;
import holeg.model.GridNode;
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) {
        List<Grid> islands = findIslands(grid);

        // Use default settings
        SolverSettings settings = new SolverSettings();

        boolean anySolved = false;
        GridSolverResult sumResult = GridSolverResult.Solved;

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

            // Try to solve problem 6 times -> each time the slack node has a new position (random)
            for (int i = 0; i < 6; i++) {
                // Build problem and solve it
                PowerFlowProblem problem = buildProblem(island);
                Solver solver = new NewtonRaphsonSolver();
                SolverResult result = solver.solve(problem, settings);

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

                // Update grid if solved
                if (result.solved) {
                    updateGrid(island, problem, result);
                    solved = true;
                    break;
                }
            }

            if (!solved)
                sumResult = GridSolverResult.Error;
            else
                anySolved = true;
        }

        // 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) {
        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 = scaleVoltages(buses);
        double scalePower = scalePower(buses);

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

        // 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 = node.getPowerConsumption().real;
        bus.Ql = node.getPowerConsumption().imaginary;
        bus.tag = node;

        if (node.getPowerGeneration().isZero()) {
            bus.Pg = 0;
            bus.Qg = 0;
            bus.type = BusType.PQ;
        }
        else {
            bus.Pg = node.getPowerGeneration().real;
            bus.Qg = node.getPowerGeneration().imaginary;
            bus.type = BusType.PV;
        }

        buses.add(bus);
    }

    private void createGridLine(Grid grid, List<Line> lines, GridEdge edge) {
        Line line = new Line();
        line.tag = edge;
        // Calculate values for line based on sample datasheet
        line.R = 0.01938; //0.791 * edge.getLengthKilometers();
        line.X = 0.1; //0.11 * edge.getLengthKilometers();
        line.B_2 = 0.1; //0.5 * 0.350141 * edge.getLengthKilometers();
        line.a = 1;
        line.from = grid.getNodes().indexOf(edge.getFrom());
        line.to = grid.getNodes().indexOf(edge.getTo());
        lines.add(line);
    }

    private double scaleVoltages(List<Bus> buses) {
        double decimal = 1;
        while(true) {
            boolean nextDecimal = false;
            for (Bus bus : buses) {
                if (bus.voltage > decimal) {
                    nextDecimal = true;
                    break;
                }
            }
            if (!nextDecimal)
                break;
            decimal *= 10;
        }
        for (Bus bus : buses)
            bus.voltage /= decimal;
        return 1.0 / decimal;
    }

    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;
        }
        return 1.0 / decimal;
    }

    private void addSlackNode(List<Bus> buses, List<Line> lines) {
        // Find possible positions
        List<Integer> possiblePositions = new ArrayList<>();
        for (int i = 0; i < buses.size(); i++)
            if (buses.get(i).Pg > 0)
                possiblePositions.add(i);

        // No generators? Add all positions
        if (possiblePositions.size() == 0) {
            for (int i = 0; i < buses.size(); i++)
                possiblePositions.add(i);
        }

        // Choose position in grid
        int position = possiblePositions.get(random.nextInt(possiblePositions.size()));

        // Create slack node
        Bus slack = new Bus();
        slack.voltage = 1;
        slack.delta = 0;
        slack.type = BusType.Slack;

        buses.add(slack);

        // Create short line
        Line line = new Line();
        line.from = position;
        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 baos = new ByteArrayOutputStream()) {
            final String utf8 = StandardCharsets.UTF_8.name();
            try (PrintStream ps = new PrintStream(baos, true, utf8)) {
                result.printTo(problem, ps);
            }
            String data = baos.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].real / problem.scaleVoltage);
            node.setCurrent(result.flowData.busCurrent[i] / currentScale);
        }

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

            edge.setPowerFlow(result.flowData.linePower[i].multiply(1 / problem.scalePower));
            edge.setLineLoss(result.flowData.linePowerLoss[i].multiply(1 / problem.scalePower));
            edge.setCurrent(result.flowData.lineCurrent[i] / currentScale);
        }
    }
}