package pl.edu.pw.mini.sg.gamegen;

import java.io.File;
import java.io.IOException;
import java.io.PrintWriter;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collections;
import java.util.HashSet;
import java.util.LinkedList;
import java.util.List;
import java.util.Optional;
import java.util.Queue;
import java.util.Set;
import java.util.function.Predicate;
import java.util.stream.Collectors;
import java.util.stream.IntStream;
import java.util.stream.Stream;

import org.apache.commons.math3.random.RandomDataGenerator;
import org.apache.commons.math3.random.RandomGenerator;

import com.codepoetics.protonpack.StreamUtils;
import com.google.common.collect.ImmutableList;

import pl.edu.pw.mini.jk.sg.game.configuration.DiscreteProbabilityFunction;
import pl.edu.pw.mini.jk.sg.game.configuration.GraphDescription;
import pl.edu.pw.mini.jk.sg.game.configuration.StaticGameConfiguration;
import pl.edu.pw.mini.jk.sg.game.graph.advanceable.AdvanceableGraphGameConfiguration;

public class BuildingLikeGameGenerator {
    private final RandomGenerator rng;
    private final BuildingLikeGameGeneratorConfig config;

    public BuildingLikeGameGenerator(RandomGenerator rng, final BuildingLikeGameGeneratorConfig config) {
        this.rng = rng;
        this.config = config;
    }

    private int[] getDistinctRadnom(int count, int upperBound) {
	RandomDataGenerator rdg = new RandomDataGenerator(rng);
        return rdg.nextPermutation(upperBound, count);
    }

    public AdvanceableGraphGameConfiguration generateRandomGame(final Optional<File> outputFilesPrefix)
            throws IOException {
        final BuildingGraphStructure structure = generateRandomStructure();
        if (outputFilesPrefix.isPresent()) {
            try (PrintWriter pw = new PrintWriter(outputFilesPrefix.get().getPath() + ".structure.txt");) {
                pw.println(structure);
            }
        } else {
            System.err.println(structure);
        }
        final Predicate<BuildingCoord> isOfficeNode = bc -> structure.getNodeRole(bc) == NodeRole.OFFICE;
        final int nodeCount = (int) structure.getAllNodes().count();
        final int officesCount = (int) structure.getAllNodes().filter(isOfficeNode).count();
        final int[] rand = getDistinctRadnom(config.getTargetCount() + 1, officesCount);
        final List<BuildingCoord> nodes = structure.getAllNodes().collect(Collectors.toList());
        final List<BuildingCoord> offices = structure.getAllNodes().filter(isOfficeNode).collect(Collectors.toList());
        final int[] randPos = Arrays.stream(rand).map(r -> nodes.indexOf(offices.get(r))).toArray();
        Collections.swap(nodes, 0, randPos[0]);

        final List<pl.edu.pw.mini.jk.sg.game.configuration.Edge> edges = new ArrayList<>();
        for (int i = 0; i < nodeCount; i++) {
            for (int j = 0; j < nodeCount; j++) {
                if (structure.areNodesConnected(nodes.get(i), nodes.get(j))) {
                    edges.add(new pl.edu.pw.mini.jk.sg.game.configuration.Edge(i, j));
                }
            }
        }

        if (outputFilesPrefix.isPresent()) {
            try (final PrintWriter pw = new PrintWriter(outputFilesPrefix.get().getPath() + ".nodepos.csv")) {
                pw.println("id,x,y,floor,type");
                StreamUtils.zipWithIndex(nodes.stream())
                        .forEach(ni -> pw.format("%d,%d,%d,%d,%s\n", ni.getIndex(), ni.getValue().getX(),
                                ni.getValue().getY(), ni.getValue().getFloor(), structure.getNodeRole(ni.getValue()).name()));
            }
        }

        final GraphDescription graph = new GraphDescription(nodeCount, edges);

        final int spawnId = (int) StreamUtils.zipWithIndex(nodes.stream())
                .filter(pair -> structure.getNodeRole(pair.getValue()) == NodeRole.ENTRANCE).findFirst().get()
                .getIndex();

        final int[][] defenderNodeRewards = IntStream
                .concat(IntStream.concat(IntStream.of(config.getDefenderSuccessRewardNormal()),
                        IntStream.generate(() -> config.getDefenderSuccessRewardTarget())
                                .limit(config.getTargetCount())),
                        IntStream.generate(() -> config.getDefenderSuccessRewardNormal()))
                .mapToObj(x -> new int[] { x }).limit(nodeCount).toArray(int[][]::new);

        final int[][] attackerNodePenalties = IntStream
                .concat(IntStream.concat(IntStream.of(config.getAttackerCaughtPenaltyNormal()),
                        IntStream.generate(() -> config.getAttackerCaughtPenaltyTarget())
                                .limit(config.getTargetCount())),
                        IntStream.generate(() -> config.getAttackerCaughtPenaltyNormal()))
                .mapToObj(x -> new int[] { x }).limit(nodeCount).toArray(int[][]::new);

        final int[][] defenderPenalties = IntStream.generate(() -> rng.nextInt(config.getDefenderPenaltyBound()))
                .limit(config.getTargetCount()).mapToObj(x -> new int[] { x }).toArray(int[][]::new);

	final int[][] attackerRewards;
	if(config.isZeroSum())
	    attackerRewards = defenderPenalties;
	else
	    attackerRewards = IntStream.generate(() -> rng.nextInt(config.getAttackerRewardBound()))
                .limit(config.getTargetCount()).mapToObj(x -> new int[] { x }).toArray(int[][]::new);

        final StaticGameConfiguration sgc = new StaticGameConfiguration(graph,
		Arrays.stream(randPos).skip(1).boxed().collect(Collectors.toList()),
                Arrays.asList(spawnId), defenderPenalties, attackerRewards, defenderNodeRewards, attackerNodePenalties,
									1, config.getNumberOfDefenders(), DiscreteProbabilityFunction.instance(), 0, 0);

        return new AdvanceableGraphGameConfiguration(sgc, config.getNumberOfAttackers());
    }

    public BuildingGraphStructure generateRandomStructure() {
        final int[] elevatorShaftX = randomShaftPositions(getWidth());
        final int[] elevatorShaftY = randomShaftPositions(getHeight());

        final List<NodeRole> nodeRoles = Stream.generate(() -> NodeRole.OFFICE)
                .limit(getWidth() * getHeight() * getFloorCount()).collect(Collectors.toList());

        final List<Boolean> xConnections = Stream.generate(() -> false)
                .limit((getWidth() - 1) * getHeight() * getFloorCount()).collect(Collectors.toList());
        final List<Boolean> yConnections = Stream.generate(() -> false)
                .limit(getWidth() * (getHeight() - 1) * getFloorCount()).collect(Collectors.toList());

        for (int f = 0; f < config.getDimensions().getFloorCount(); f++) {
            final int fl = f;
            for (int e = 0; e < config.getElevatorCount(); e++) {
                final int x = elevatorShaftX[e];
                final int y = elevatorShaftY[e];
                Stream.concat(
                        IntStream.range(0, getWidth()).mapToObj(
                                xx -> BuildingCoordsHelpers.convert3dCoordsToIndex(xx, y, fl, config.getDimensions())),
                        IntStream.range(0, getHeight()).mapToObj(
                                yy -> BuildingCoordsHelpers.convert3dCoordsToIndex(x, yy, fl, config.getDimensions())))
                        .forEach(idx -> nodeRoles.set(idx, NodeRole.CORRIDOR));
                nodeRoles.set(BuildingCoordsHelpers.convert3dCoordsToIndex(x, y, f, config.getDimensions()),
                        NodeRole.ELEVATOR);
            }

            generateFloor(f, nodeRoles, xConnections, yConnections);
        }

        final int edge;
        if (elevatorShaftX[0] == 0) {
            edge = getWidth() - 1;
        } else {
            edge = 0;
        }

        nodeRoles.set(BuildingCoordsHelpers.convert3dCoordsToIndex(edge, elevatorShaftY[0], 0, config.getDimensions()),
                NodeRole.ENTRANCE);

        return new BuildingGraphStructure(ImmutableList.copyOf(nodeRoles), config.getDimensions(),
                ImmutableList.copyOf(xConnections), ImmutableList.copyOf(yConnections));
    }

    private void generateFloor(final int f, final List<NodeRole> nodeRoles, final List<Boolean> xConnections,
            final List<Boolean> yConnections) {
        ImmutableList<Door> doors = addMissingDoors(f, nodeRoles, generateDoors(f, nodeRoles));
        for (int x = 0; x < config.getDimensions().getWidth(); x++) {
            for (int y = 0; y < config.getDimensions().getHeight(); y++) {
                final BuildingCoord p = new BuildingCoord(x, y, f);
                final Optional<BuildingCoord> nX = p.advanceX(config.getDimensions());
                final Optional<BuildingCoord> nY = p.advanceY(config.getDimensions());
                if (nX.isPresent() && isMainCorridor(nodeRoles, p, nX.get())) {
                    xConnections.set(BuildingCoordsHelpers.getXConnectionIndex(p.getX(), nX.get().getX(), p.getY(),
                            p.getFloor(), config.getDimensions()), true);
                }
                if (nY.isPresent() && isMainCorridor(nodeRoles, p, nY.get())) {
                    yConnections.set(BuildingCoordsHelpers.getYConnectionIndex(p.getX(), p.getY(), nY.get().getY(),
                            p.getFloor(), config.getDimensions()), true);
                }
            }
        }

        doors.stream().forEach((door) -> addEdge(door.getCorridor(), door.getOffice(), xConnections, yConnections));

        buildOfficeTrees(f, nodeRoles, doors, xConnections, yConnections);
        buildRandomOfficeEdges(f, nodeRoles, xConnections, yConnections);
    }

    private void buildRandomOfficeEdges(final int floor, final List<NodeRole> nodeRoles,
            final List<Boolean> xConnections, final List<Boolean> yConnections) {
        BuildingCoordsHelpers.nodesOnFloor(floor, config.getDimensions())
                .filter(bc -> nodeRoles.get(bc.to3dIndex(config.getDimensions())) == NodeRole.OFFICE).forEach((bc) -> {
                    final Optional<BuildingCoord> nX = bc.advanceX(config.getDimensions());
                    final Optional<BuildingCoord> nY = bc.advanceY(config.getDimensions());

                    if (nX.isPresent() && nodeRoles.get(nX.get().to3dIndex(config.getDimensions())) == NodeRole.OFFICE
                            && rng.nextDouble() < config.getInteriorDensity()) {
                        addEdge(bc, nX.get(), xConnections, yConnections);
                    }

                    if (nY.isPresent() && nodeRoles.get(nY.get().to3dIndex(config.getDimensions())) == NodeRole.OFFICE
                            && rng.nextDouble() < config.getInteriorDensity()) {
                        addEdge(bc, nY.get(), xConnections, yConnections);
                    }

                });
    }

    private void buildOfficeTrees(final int floor, final List<NodeRole> nodeRoles, final ImmutableList<Door> doors,
            final List<Boolean> xConnections, final List<Boolean> yConnections) {
        final Set<BuildingCoord> connected = new HashSet<>();
        final Set<BuildingCoord> notConnected = BuildingCoordsHelpers.nodesOnFloor(floor, config.getDimensions())
                .filter(bc -> nodeRoles.get(bc.to3dIndex(config.getDimensions())) == NodeRole.OFFICE)
                .collect(Collectors.toSet());

        for (final Door e : doors) {
            final BuildingCoord coord = e.getOffice();
            notConnected.remove(coord);
            connected.add(coord);
        }

        while (!notConnected.isEmpty()) {
            final BuildingCoord randomConnected = connected.stream().skip(rng.nextInt(connected.size())).findFirst()
                    .get();
            final List<BuildingCoord> neighbors = randomConnected.getNeighbours(config.getDimensions()).stream()
                    .filter(notConnected::contains).collect(Collectors.toList());
            if (neighbors.size() == 0) {
                connected.remove(randomConnected);
            } else if (neighbors.size() == 1) {
                final BuildingCoord newConnected = neighbors.get(0);
                connected.remove(randomConnected);
                notConnected.remove(newConnected);
                connected.add(newConnected);
                addEdge(randomConnected, newConnected, xConnections, yConnections);
            } else {
                final BuildingCoord newConnected = neighbors.get(rng.nextInt(neighbors.size()));
                notConnected.remove(newConnected);
                connected.add(newConnected);
                addEdge(randomConnected, newConnected, xConnections, yConnections);
            }
        }
    }

    private void addEdge(BuildingCoord n1, BuildingCoord n2, List<Boolean> xConnections, List<Boolean> yConnections) {
        if (BuildingCoord.areXNeighbours(n1, n2)) {
            xConnections.set(BuildingCoordsHelpers.getXConnectionIndex(n1.getX(), n2.getX(), n1.getY(), n1.getFloor(),
                    config.getDimensions()), true);
        } else if (BuildingCoord.areYNeighbours(n1, n2)) {
            yConnections.set(BuildingCoordsHelpers.getYConnectionIndex(n1.getX(), n1.getY(), n2.getY(), n1.getFloor(),
                    config.getDimensions()), true);
        } else {
            throw new IllegalArgumentException();
        }
    }

    private boolean isMainCorridor(final List<NodeRole> nodeRoles, final BuildingCoord p1, final BuildingCoord p2) {
        final NodeRole r1 = nodeRoles.get(BuildingCoordsHelpers.convert3dCoordsToIndex(p1, config.getDimensions()));
        final NodeRole r2 = nodeRoles.get(BuildingCoordsHelpers.convert3dCoordsToIndex(p2, config.getDimensions()));

        return r1.isMainRoute() && r2.isMainRoute();
    }

    private ImmutableList<Door> generateDoors(final int floor, final List<NodeRole> nodeRoles) {
        ImmutableList.Builder<Door> builder = ImmutableList.builder();
        for (int x = 0; x < getWidth(); x++) {
            for (int y = 0; y < getWidth(); y++) {
                final BuildingCoord p1 = new BuildingCoord(x, y, floor);
                for (BuildingCoord p2 : p1.getNeighbours(config.getDimensions())) {
                    if (nodeRoles.get(p1.to3dIndex(config.getDimensions())).isMainRoute()
                            && nodeRoles.get(p2.to3dIndex(config.getDimensions())) == NodeRole.OFFICE
                            && rng.nextDouble() < config.getDoorDensity()) {
                        builder.add(new Door(p1, p2));
                    }
                }
            }
        }

        return builder.build();
    }

    private ImmutableList<Door> addMissingDoors(final int floor, final List<NodeRole> nodeRoles,
            final ImmutableList<Door> randomDoors) {
        final Set<BuildingCoord> unreachable = officesSet(floor, nodeRoles);
        final ImmutableList.Builder<Door> builder = ImmutableList.builder();
        builder.addAll(randomDoors);

        for (final Door e : randomDoors) {
            if (unreachable.contains(e.getOffice())) {
                unreachable.removeAll(officeSubset(nodeRoles, e.getOffice()));
            }
        }

        while (!unreachable.isEmpty()) {
            final BuildingCoord possibleDoor = unreachable.stream().skip(rng.nextInt(unreachable.size())).findFirst()
                    .get();
            final List<BuildingCoord> possibleDoorOutputs = possibleDoor.getNeighbours(config.getDimensions()).stream()
                    .filter(bc -> nodeRoles.get(bc.to3dIndex(config.getDimensions())).isMainRoute())
                    .collect(Collectors.toList());
            if (possibleDoorOutputs.size() > 0) {
                final BuildingCoord doorOutput = possibleDoorOutputs.get(rng.nextInt(possibleDoorOutputs.size()));
                unreachable.removeAll(officeSubset(nodeRoles, possibleDoor));
                builder.add(new Door(doorOutput, possibleDoor));
            }
        }

        return builder.build();
    }

    private Set<BuildingCoord> officesSet(int floor, List<NodeRole> nodeRoles) {
        final HashSet<BuildingCoord> set = new HashSet<>();
        for (int x = 0; x < config.getDimensions().getWidth(); x++) {
            for (int y = 0; y < config.getDimensions().getHeight(); y++) {
                final BuildingCoord c = new BuildingCoord(x, y, floor);
                if (nodeRoles.get(c.to3dIndex(config.getDimensions())) == NodeRole.OFFICE) {
                    set.add(c);
                }
            }
        }
        return set;
    }

    private Set<BuildingCoord> officeSubset(final List<NodeRole> nodeRoles, final BuildingCoord startPoint) {
        if (nodeRoles.get(startPoint.to3dIndex(config.getDimensions())) != NodeRole.OFFICE) {
            throw new IllegalArgumentException();
        }
        HashSet<BuildingCoord> reachable = new HashSet<>(Arrays.asList(startPoint));
        Queue<BuildingCoord> nodesToProcess = new LinkedList<>(Arrays.asList(startPoint));

        while (!nodesToProcess.isEmpty()) {
            final BuildingCoord currentNode = nodesToProcess.poll();
            final Set<BuildingCoord> interestingNeighbours = currentNode.getNeighbours(config.getDimensions()).stream()
                    .filter(bc -> NodeRole.OFFICE == nodeRoles.get(bc.to3dIndex(config.getDimensions())))
                    .filter(((Predicate<BuildingCoord>) reachable::contains).negate()).collect(Collectors.toSet());
            reachable.addAll(interestingNeighbours);
            nodesToProcess.addAll(interestingNeighbours);
        }

        return reachable;
    }

    private int[] randomShaftPositions(int limit) {
        int[] ret;
        do {
            ret = IntStream.generate(() -> rng.nextInt(limit)).limit(config.getElevatorCount()).toArray();
        } while (!checkCoordPositions(ret));
        return ret;
    }

    private boolean checkCoordPositions(final int[] x) {
        int[] sortedX = Arrays.stream(x).sorted().toArray();
        return StreamUtils
                .zip(Arrays.stream(sortedX).boxed(), Arrays.stream(sortedX).skip(1).boxed(), (z, y) -> (y - z))
                .allMatch(i -> i > 1);
    }

    public int getWidth() {
        return config.getDimensions().getWidth();
    }

    public int getHeight() {
        return config.getDimensions().getHeight();
    }

    public int getFloorCount() {
        return config.getDimensions().getFloorCount();
    }

}
