package main.java.com.example;

import org.apache.ignite.Ignite;
import org.apache.ignite.Ignition;
import org.apache.ignite.compute.ComputeJob;
import org.apache.ignite.compute.ComputeJobResult;
import org.apache.ignite.compute.ComputeTaskAdapter;
import org.apache.ignite.cluster.ClusterNode;

import java.io.InputStream;
import java.net.URL;
import java.util.*;

public class MainTask {

    public static void main(String[] args) {
        ClassLoader classLoader = MainTask.class.getClassLoader();
        URL resource = classLoader.getResource("ignite.xml");

        if (resource == null) {
            System.out.println("ignite.xml not found in classpath!");
            return;
        }

        try (InputStream igniteConfigStream = classLoader.getResourceAsStream("ignite.xml")) {
            if (igniteConfigStream == null) {
                System.out.println("Failed to load ignite.xml from resources.");
                return;
            }

            try (Ignite ignite = Ignition.start(igniteConfigStream)) {
                System.out.println("Ignite started successfully.");

                int rows = 3000;
                int cols = 3000;
                int[][] matrix = generateRandomMatrix(rows, cols);

                System.out.println("Original Matrix:");
                printMatrix(matrix);

                long startTime = System.nanoTime();
                try {
                    int[][] sortedMatrix = ignite.compute().execute(new MatrixSortingTask(), matrix);
                    long endTime = System.nanoTime();

                    System.out.println("Sorted Matrix:");
                    printMatrix(sortedMatrix);
                    System.out.println("Time taken: " + (endTime - startTime) / 1_000_000 + " ms");
                } catch (Exception e) {
                    System.out.println("Error during matrix sorting computation: " + e.getMessage());
                    e.printStackTrace();
                }
            }
        } catch (Exception e) {
            System.out.println("Ignite failed to start: " + e.getMessage());
            e.printStackTrace();
        }
    }

    public static int[][] generateRandomMatrix(int rows, int cols) {
        Random random = new Random();
        int[][] matrix = new int[rows][cols];
        for (int i = 0; i < rows; i++) {
            for (int j = 0; j < cols; j++) {
                matrix[i][j] = random.nextInt(100);
            }
        }
        System.out.println("Random matrix generated.");
        return matrix;
    }

    public static void printMatrix(int[][] matrix) {
        for (int i = 0; i < 10 && i < matrix.length; i++) {
            for (int j = 0; j < 10 && j < matrix[0].length; j++) {
                System.out.print(matrix[i][j] + " ");
            }
            System.out.println();
        }
    }

    private static class MatrixSortingTask extends ComputeTaskAdapter<int[][], int[][]> {

        @Override
        public Map<? extends ComputeJob, ClusterNode> map(List<ClusterNode> nodes, int[][] matrix) {
            System.out.println("Mapping jobs to nodes. Total nodes: " + nodes.size());
            Map<ComputeJob, ClusterNode> map = new HashMap<>();
            int cols = matrix[0].length;
            int nodeCount = nodes.size();
            int chunkSize = (int) Math.ceil((double) cols / nodeCount);

            for (int i = 0; i < nodeCount; i++) {
                int startCol = i * chunkSize;
                int endCol = Math.min(startCol + chunkSize, cols);

                if (startCol >= endCol) break;

                int[][] subMatrix = extractColumns(matrix, startCol, endCol);
                map.put(new ColumnSortJob(subMatrix, startCol), nodes.get(i));
            }
            return map;
        }

        @Override
        public int[][] reduce(List<ComputeJobResult> results) {
            System.out.println("Reducing results from jobs. Number of results: " + results.size());
            if (results.isEmpty()) {
                return new int[0][0];
            }

            int[][] firstChunk = (int[][]) results.get(0).getData();
            int rows = firstChunk.length;
            int cols = results.stream().mapToInt(r -> ((int[][]) r.getData())[0].length).sum();
            int[][] sortedMatrix = new int[rows][cols];

            int currentCol = 0;
            for (ComputeJobResult result : results) {
                int[][] sortedChunk = (int[][]) result.getData();
                mergeColumns(sortedMatrix, sortedChunk, currentCol);
                currentCol += sortedChunk[0].length;
            }
            return sortedMatrix;
        }

        private int[][] extractColumns(int[][] matrix, int startCol, int endCol) {
            int rows = matrix.length;
            int[][] subMatrix = new int[rows][endCol - startCol];
            for (int i = 0; i < rows; i++) {
                System.arraycopy(matrix[i], startCol, subMatrix[i], 0, endCol - startCol);
            }
            return subMatrix;
        }

        private void mergeColumns(int[][] target, int[][] sortedChunk, int startCol) {
            for (int i = 0; i < target.length; i++) {
                System.arraycopy(sortedChunk[i], 0, target[i], startCol, sortedChunk[i].length);
            }
        }
    }

    private static class ColumnSortJob implements ComputeJob {
        private final int[][] matrix;
        private final int startCol;

        public ColumnSortJob(int[][] matrix, int startCol) {
            this.matrix = matrix;
            this.startCol = startCol;
        }

        @Override
        public Object execute() {
            System.out.println("Sorting columns starting from index " + startCol);
            int rows = matrix.length;
            int cols = matrix[0].length;
            Integer[] indices = new Integer[cols];
            for (int i = 0; i < cols; i++) indices[i] = i;
            Arrays.sort(indices, Comparator.comparingInt(a -> matrix[0][a]));

            int[][] sortedMatrix = new int[rows][cols];
            for (int i = 0; i < cols; i++) {
                for (int j = 0; j < rows; j++) {
                    sortedMatrix[j][i] = matrix[j][indices[i]];
                }
            }
            System.out.println("Sorting complete for columns starting at index " + startCol);
            return sortedMatrix;
        }

        @Override
        public void cancel() {
            System.out.println("ColumnSortJob cancelled.");
        }
    }
}