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.*;
import java.util.concurrent.Semaphore;

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

    public GridSolverResult solve(Grid grid, PowerFlowSettings settings, HolegPowerFlowContext context) {
        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) {
            // Check if we should stop
            if (Thread.interrupted())
                return GridSolverResult.Interrupted;

            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;

            // Build all problems
            PowerFlowProblem[] problems = new PowerFlowProblem[solveTryCount];
            SolverResult[] results = new SolverResult[solveTryCount];

            for (int i = 0; i < solveTryCount; i++) {
                // Check if we should stop
                if (Thread.interrupted())
                    return GridSolverResult.Interrupted;

                // 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;
                }

                problems[i] = problem;
                results[i] = null;
            }

            // Create solver jobs
            if (!solved) {
                Semaphore semaphore = new Semaphore(0);
                Thread[] jobs = new Thread[problems.length];
                for (int i = 0; i < problems.length; i++) {
                    int jobID = i;
                    PowerFlowProblem problem = problems[jobID];
                    jobs[jobID] = new Thread(() -> {
                        Solver solver = new NewtonRaphsonSolver();
                        SolverResult result = solver.solve(problem, settings.solverSettings);

                        synchronized (results) {
                            results[jobID] = result;
                        }

                        // Notify main thread that we are done
                        semaphore.release();
                    });
                    jobs[jobID].setName("Solver Job " + Long.toHexString(jobs[jobID].getId()));
                }

                // Check if we should have stopped
                if (Thread.interrupted())
                    return GridSolverResult.Interrupted;

                System.out.println("     ");

                // Check if actually we have jobs to run
                if (jobs.length > 0) {
                    boolean[] resultUsed = new boolean[jobs.length];

                    // Start maximum number of threads
                    int jobThreadCount = Math.min(jobs.length, Math.max(1, settings.maxSolverThreads));
                    for (int i = 0; i < jobThreadCount; i++)
                        jobs[i].start();

                    int nextJob = jobThreadCount;
                    while (true) {
                        // Wait for any job to finish
                        try {
                            semaphore.acquire();
                        } catch (InterruptedException ignored) {
                            for (Thread job : jobs) job.interrupt();
                            return GridSolverResult.Interrupted;
                        }

                        // Check results of all jobs
                        boolean allSolved = true;

                        synchronized (results) {
                            for (int i = 0; i < jobs.length; i++) {
                                // Already checked?
                                if (resultUsed[i])
                                    continue;

                                // Check result of the job
                                PowerFlowProblem problem = problems[i];
                                SolverResult result = results[i];
                                if (result != null) {
                                    // Debug message
                                    if (PowerFlowAnalysisMenu.getInstance().shouldShowDebug() && (result.solved || i == solveTryCount - 1))
                                        debugShowResultAsMessageBox(problem, result);

                                    // Solver was interrupted
                                    if (result.error == SolverError.Interrupted)
                                        return GridSolverResult.Interrupted;

                                    System.out.printf("Problem %d, solved: %s, slack: %.3f\n", i, result.solved ? "true" : "false", result.solved ? result.flowData.busInjection[i].real / problem.scalePower : 0);

                                    // Solver was successful?
                                    if (result.solved) {
                                        // Simply use the first solved result
                                        if (settings.slackNodePlacementStrategy != SlackNodePlacementStrategy.MinimizeSlack) {
                                            solved = true;
                                            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;
                                        }
                                    }
                                    resultUsed[i] = true;

                                } else
                                    allSolved = false;
                            }
                        }

                        // Either all are solved or we solved one job
                        if (solved || allSolved)
                            break;

                        // Check if we should have stopped
                        if (Thread.interrupted()) {
                            for (Thread job : jobs) job.interrupt();
                            return GridSolverResult.Interrupted;
                        }

                        // Start next thread
                        if (nextJob < jobs.length)
                            jobs[nextJob++].start();
                    }
                    // Stop every job
                    for (Thread job : jobs) job.interrupt();
                }
            }

            // Update grid if solved
            if (solved || bestResult != null) {
                // 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);

                        // Save problem and result to context
                        if (context != null) {
                            synchronized (context.lock) {
                                context.problem = bestProblem;
                                context.result = bestResult;
                            }
                        }
                    }
                    anySolved = true;
                }
                else
                    sumResult = GridSolverResult.Error;
            }
            else
                sumResult = GridSolverResult.Error;
        }

        // Check if we should have stopped
        if (Thread.interrupted())
            return GridSolverResult.Interrupted;

        // 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);

        // Scale voltages and scale power
        double scaleVoltage = settings.disableAutomaticScaling ? 1 : 1.0 / 230.0;
        double scalePower = settings.disableAutomaticScaling ? 1 : scalePower(buses);
        double scaleResistance = settings.disableAutomaticScaling ? 1 :  1.0 / ((scaleVoltage * scaleVoltage) / scalePower);

        for (GridEdge edge : grid.getEdges())
            createGridLine(grid, lines, edge, scaleResistance);

        // 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.Bus) {
            switch(node.getTypeByDesign()) {
                case PV:
                    bus.type = BusType.PV;
                    break;
                case PQ:
                    bus.type = BusType.PQ;
                    break;
                case Slack:
                    bus.type = BusType.Slack;
                    break;
            }
            bus.Pg = Math.abs(node.getPowerGeneration().real);
            bus.Qg = Math.abs(node.getPowerGeneration().imaginary);
        }
        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;
        bus.index = buses.size();

        buses.add(bus);
    }

    private void createGridLine(Grid grid, List<Line> lines, GridEdge edge, double scaleResistance) {
        Line line = new Line();
        line.tag = edge;
        if (ComplexNumber.isNullOrZero(edge.getOverrideImpedance())) {
            // Calculate values for line based on sample datasheet
            line.R = 0.00642 * edge.getLengthKilometers() * scaleResistance;
            line.X = 0.0083 * edge.getLengthKilometers() * scaleResistance;
            line.B_2 = (0.0001 * edge.getLengthKilometers()) / scaleResistance;
        }
        else {
            line.R = edge.getOverrideImpedance().real;
            line.X = edge.getOverrideImpedance().imaginary;
            line.B_2 = edge.getOverrideShuntSusceptance();
            if (!edge.getOverrideInPerUnit()) {
                line.R *= scaleResistance;
                line.X *= scaleResistance;
                line.B_2 *= scaleResistance;
            }
        }
        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);

            node.setIndex(problem.buses[i].index);

            // 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);
        }
    }
}