package org.kne.codec.kif;

import javax.imageio.ImageIO;
import javax.imageio.ImageReader;
import javax.imageio.ImageWriter;
import javax.imageio.spi.IIORegistry;
import javax.imageio.spi.ImageReaderSpi;
import javax.imageio.spi.ImageWriterSpi;
import java.awt.image.BufferedImage;
import java.io.*;
import java.nio.charset.StandardCharsets;
import java.nio.file.*;
import java.util.*;
import java.util.stream.Collectors;

/**
 * 图像压缩格式对比实验
 * 遍历 input 目录下的所有图片，用 ImageIO 注册的所有格式进行压缩，
 * 输出压缩后大小、压缩率，并验证无损/有损。
 * 分别记录编码时间和解码时间
 */
public class CompressionExperiment {

    // 支持的输入格式（ImageIO 能读取的）
    private static final Set<String> INPUT_FORMATS = new HashSet<>(Arrays.asList(
        "png", "bmp", "kif"
    ));
    // 要测试的输出格式（从 ImageIO 注册的 Writer 中获取）
    private static Set<String> outputFormats=new HashSet<>(Arrays.asList(
            "png", "kif"
    	    ));

    public static void main(String[] args) throws Exception {
    	 // 🔥 强制覆盖标准输出流为 UTF-8
        System.setOut(new PrintStream(System.out, true, StandardCharsets.UTF_8));
        System.setErr(new PrintStream(System.err, true, StandardCharsets.UTF_8));
        
        // 1. 注册 KIF 插件
        org.kne.codec.kif.KIFImageReaderSpi.register();
        org.kne.codec.kif.KIFImageWriterSpi.register();

        // 2. 获取所有可用的输出格式
        System.out.println("=== 可用的输出格式 ===");
        outputFormats.forEach(f -> System.out.println("  " + f));
        System.out.println();

        // 3. 准备目录
        Path inputDir = Paths.get("input2");
        Path outputDir = Paths.get("output");
        if (!Files.exists(inputDir)) {
            System.err.println("错误: input 目录不存在！");
            return;
        }
        if (!Files.exists(outputDir)) {
            Files.createDirectories(outputDir);
        }

        // 4. 获取所有测试图片
        List<Path> imageFiles = getTestImages(inputDir);
        if (imageFiles.isEmpty()) {
            System.err.println("错误: input 目录中没有找到图片！");
            return;
        }
        System.out.println("找到 " + imageFiles.size() + " 张测试图片\n");
        // 5.预热
        warmup(outputDir, imageFiles,30000000000L);
        // 6. 执行实验
        List<ExperimentResult> results = runExperiment(outputDir, imageFiles);

        // 7. 输出汇总报告
        printSummary(results);
        
        System.out.println(KIFImageWriter.getModeCollector());
    }
    private static void warmup(Path outputDir, List<Path> imageFiles, long warmupTimeNs) {
        long startTime = System.nanoTime();
        long elapsed = 0;
        int iteration = 0;
        
        System.out.println("========================================");
        System.out.println("🔥 预热开始 (目标: " + (warmupTimeNs / 1_000_000) + "ms)");
        System.out.println("========================================");
        
            loop:while(true) {
            for (Path imageFile : imageFiles) {
            	if(elapsed >= warmupTimeNs) {
            		break loop;
            	}
            long iterStart = System.nanoTime();
                try {
                    BufferedImage original = ImageIO.read(imageFile.toFile());
                    if (original == null) {
                        continue;
                    }
                    
                    // 对所有格式进行预热压缩和解压（丢弃结果）
                    for (String format : outputFormats) {
                        try {
                            // 执行完整的编码+解码流程，但不保存结果
                            Path tempFile = outputDir.resolve("warmup_" + iteration + "_" + 
                                imageFile.getFileName().toString().replaceAll("\\.[^.]*$", "") + 
                                "." + format.toLowerCase());
                            
                            // 编码
                            boolean written = ImageIO.write(original, format, tempFile.toFile());
                            if (!written) {
                                continue;
                            }
                            
                            // 解码（验证无损性）
                            BufferedImage decoded = ImageIO.read(tempFile.toFile());
                            if (decoded != null) {
                                // 简单比较尺寸，不逐像素比较（节省时间）
                                if (decoded.getWidth() == original.getWidth() && 
                                    decoded.getHeight() == original.getHeight()) {
                                    // 预热成功
                                }
                            }
                            
                            // 立即删除临时文件
                            Files.deleteIfExists(tempFile);
                            
                        } catch (Exception e) {
                            // 预热时忽略错误
                        }
                    }
                    
                } catch (Exception e) {
                    // 忽略预热错误
                }
            iteration++;
            long iterEnd = System.nanoTime();
            elapsed = iterEnd - startTime;
            }
            }
            
           
        
        
        System.out.println("========================================");
        System.out.printf("🔥 预热完成: %d 张图片, 耗时 %dms\n", 
            iteration, elapsed / 1_000_000L);
        System.out.println("========================================");
        
        // 清理临时文件
        try {
            Files.list(outputDir)
                .filter(p -> p.getFileName().toString().startsWith("warmup_"))
                .forEach(p -> {
                    try { Files.deleteIfExists(p); } catch (Exception e) {}
                });
        } catch (Exception e) {
            // 忽略清理错误
        }
    }
	private static List<ExperimentResult> runExperiment(Path outputDir, List<Path> imageFiles) {
		List<ExperimentResult> results = new ArrayList<>();
        for (Path imageFile : imageFiles) {
            try {
            	 // 读取原图
                long readStart = System.nanoTime();
                BufferedImage original = ImageIO.read(imageFile.toFile());
                long readEnd = System.nanoTime();
                long readTimeMs = (readEnd - readStart) / 1_000_000;
                
                if (original == null) {
                    System.err.println("  ⚠️ 无法读取:"+imageFile.getFileName()+"，跳过\n");
                    continue;
                }
                
            System.out.println("========================================");
            System.out.println("测试图片: " + imageFile.getFileName()+" ("+original.getWidth()+"x"+original.getHeight()+")");
            System.out.println("========================================");
            
               
                
                long originalSize = Files.size(imageFile);
                System.out.printf("  原始大小: %,d bytes (%.2f KB)  | 读取耗时: %dms\n", 
                    originalSize, originalSize / 1024.0, readTimeMs);
                
                // 对每种输出格式进行测试
                for (String format : outputFormats) {
                    try {
                        ExperimentResult result = testFormat(original, imageFile, format, outputDir);
                        results.add(result);
                        System.out.println(result);
                    } catch (Exception e) {
                        System.err.println("  ❌ " + format + " 测试失败: " );
                        e.printStackTrace();
                    }
                }
                System.out.println();
                
            } catch (Exception e) {
                System.err.println("  ❌ 处理 " + imageFile.getFileName() + " 时出错: " );
                e.printStackTrace();
            }
        }
		return results;
	}

    /**
     * 获取所有可用的图像写入格式
     */
    private static Set<String> getAvailableWriterFormats() {
        Set<String> formats = new HashSet<>();
        IIORegistry registry = IIORegistry.getDefaultInstance();
        Iterator<ImageWriterSpi> iter = registry.getServiceProviders(ImageWriterSpi.class, true);
        while (iter.hasNext()) {
            ImageWriterSpi spi = iter.next();
            String[] names = spi.getFormatNames();
            if (names != null) {
                for (String name : names) {
                    formats.add(name.toLowerCase());
                }
            }
        }
        return formats;
    }

    /**
     * 获取 input 目录下的所有图片文件
     */
    private static List<Path> getTestImages(Path inputDir) throws IOException {
        List<Path> images = new ArrayList<>();
        try (DirectoryStream<Path> stream = Files.newDirectoryStream(inputDir)) {
            for (Path entry : stream) {
                if (Files.isRegularFile(entry)) {
                    String fileName = entry.getFileName().toString().toLowerCase();
                    for (String fmt : INPUT_FORMATS) {
                        if (fileName.endsWith("." + fmt)) {
                            images.add(entry);
                            break;
                        }
                    }
                }
            }
        }
        images.sort(Comparator.comparing(p -> p.getFileName().toString()));
        return images;
    }

    /**
     * 测试一种格式的压缩效果（纯内存测试，排除硬盘 I/O）
     * 分别记录编码时间和解码时间
     */
    private static ExperimentResult testFormat(BufferedImage original, Path sourceFile, 
                                                String format, Path outputDir) throws Exception {
        // 生成输出文件名（仅用于显示和保存验证）
        String baseName = sourceFile.getFileName().toString();
        int dotIdx = baseName.lastIndexOf('.');
        String nameWithoutExt = (dotIdx > 0) ? baseName.substring(0, dotIdx) : baseName;
        String outputFileName = nameWithoutExt + "." + format.toLowerCase();
        Path outputFile = outputDir.resolve(outputFileName);
        
        // ===== 编码时间（纯内存） =====
        long encodeStart = System.nanoTime();
        
        // 使用 ByteArrayOutputStream 替代 FileOutputStream
        int estimatedSize = (int) (original.getWidth()*original.getHeight()*4 * 1.5);
        ByteArrayOutputStream baos = new ByteArrayOutputStream(estimatedSize);
        boolean written = ImageIO.write(original, format, baos);
        
        long encodeEnd = System.nanoTime();
        long encodeTimeMs = (encodeEnd - encodeStart) / 1_000_000;
        
        if (!written) {
            throw new IOException("ImageIO.write() 返回 false，可能不支持此格式");
        }
        
        // 压缩后大小（从内存中获取）
        byte[] encodedData = baos.toByteArray();
        long compressedSize = encodedData.length;
        long originalSize = Files.size(sourceFile);
        double compressionRatio = (double) compressedSize / originalSize * 100;
        
        // ===== 解码时间（纯内存） =====
        long decodeStart = System.nanoTime();
        
        // 使用 ByteArrayInputStream 替代 FileInputStream
        ByteArrayInputStream bais = new ByteArrayInputStream(encodedData);
        BufferedImage decoded = ImageIO.read(bais);
        
        long decodeEnd = System.nanoTime();
        long decodeTimeMs = (decodeEnd - decodeStart) / 1_000_000;
        
        // 验证无损性
        boolean lossless = compareImages(original, decoded);
        
        // ===== 可选：将结果保存到磁盘（用于调试/验证） =====
        // 如果需要实际保存文件以便后续查看，取消注释下面的代码
        // try (FileOutputStream fos = new FileOutputStream(outputFile.toFile())) {
        //     fos.write(encodedData);
        // }
        
        return new ExperimentResult(
            sourceFile.getFileName().toString(),
            format,
            originalSize,
            compressedSize,
            compressionRatio,
            encodeTimeMs,
            decodeTimeMs,
            lossless
        );
    }

    /**
     * 逐像素比较两张图片是否完全相同
     */
    private static boolean compareImages(BufferedImage img1, BufferedImage img2) {
        if (img1 == null || img2 == null) return false;
        if (img1.getWidth() != img2.getWidth()) return false;
        if (img1.getHeight() != img2.getHeight()) return false;
        
        int width = img1.getWidth();
        int height = img1.getHeight();
        
        for (int y = 0; y < height; y++) {
            for (int x = 0; x < width; x++) {
                if (img1.getRGB(x, y) != img2.getRGB(x, y)) {
                    return false;
                }
            }
        }
        return true;
    }

    /**
     * 打印汇总报告
     */
    private static void printSummary(List<ExperimentResult> results) {
        System.out.println("\n\n");
        System.out.println("╔═══════════════════════════════════════════════════════════════════════════════════════════════════╗");
        System.out.println("║                                    汇总报告                                                    ║");
        System.out.println("╚═══════════════════════════════════════════════════════════════════════════════════════════════════╝");
        
        // 按格式分组统计
        Map<String, List<ExperimentResult>> grouped = results.stream()
            .collect(Collectors.groupingBy(r -> r.format));
        
        // 表头 - 压缩率统计
        System.out.println("\n┌──────────┬────────────┬──────────────┬─────────────┬────────────┬──────────┐");
        System.out.println("│  格式    │  测试数量  │  平均压缩率  │  最大压缩率  │  最小压缩率 │  无损率  │");
        System.out.println("├──────────┼────────────┼──────────────┼─────────────┼────────────┼──────────┤");
        
        for (Map.Entry<String, List<ExperimentResult>> entry : grouped.entrySet()) {
            String format = entry.getKey();
            List<ExperimentResult> list = entry.getValue();
            
            double avgRatio = list.stream().mapToDouble(r -> r.compressionRatio).average().orElse(0);
            double maxRatio = list.stream().mapToDouble(r -> r.compressionRatio).max().orElse(0);
            double minRatio = list.stream().mapToDouble(r -> r.compressionRatio).min().orElse(0);
            long losslessCount = list.stream().filter(r -> r.lossless).count();
            double losslessRate = (double) losslessCount / list.size() * 100;
            
            System.out.printf("│  %-8s│  %10d│  %12.2f%%│  %11.2f%%│  %10.2f%%│  %8.1f%%│\n",
                format, list.size(), avgRatio, maxRatio, minRatio, losslessRate);
        }
        System.out.println("└──────────┴────────────┴──────────────┴─────────────┴────────────┴──────────┘");
        
        // 表头 - 时间统计
        System.out.println("\n┌──────────┬────────────┬──────────────┬──────────────┬──────────────────┐");
        System.out.println("│  格式    │  测试数量  │  平均编码时间 │  平均解码时间 │  编码/解码比     │");
        System.out.println("├──────────┼────────────┼──────────────┼──────────────┼──────────────────┤");
        
        for (Map.Entry<String, List<ExperimentResult>> entry : grouped.entrySet()) {
            String format = entry.getKey();
            List<ExperimentResult> list = entry.getValue();
            
            double avgEncode = list.stream().mapToLong(r -> r.encodeTimeMs).average().orElse(0);
            double avgDecode = list.stream().mapToLong(r -> r.decodeTimeMs).average().orElse(0);
            double ratio = avgDecode > 0 ? avgEncode / avgDecode : 0;
            
            System.out.printf("│  %-8s│  %10d│  %12.1fms│  %12.1fms│  %8.2fx     │\n",
                format, list.size(), avgEncode, avgDecode, ratio);
        }
        System.out.println("└──────────┴────────────┴──────────────┴──────────────┴──────────────────┘");
        
        // KIF 专项统计
        List<ExperimentResult> kifResults = results.stream()
            .filter(r -> "kif".equalsIgnoreCase(r.format))
            .collect(Collectors.toList());
        
        if (!kifResults.isEmpty()) {
            System.out.println("\n📊 KIF 格式详细统计:");
            System.out.println("  测试图片数: " + kifResults.size());
            double avgKifRatio = kifResults.stream().mapToDouble(r -> r.compressionRatio).average().orElse(0);
            System.out.printf("  平均压缩率: %.2f%% (相对于原始文件)\n", avgKifRatio);
            long losslessKif = kifResults.stream().filter(r -> r.lossless).count();
            System.out.println("  无损图片数: " + losslessKif + "/" + kifResults.size());
            
            double avgKifEncode = kifResults.stream().mapToLong(r -> r.encodeTimeMs).average().orElse(0);
            double avgKifDecode = kifResults.stream().mapToLong(r -> r.decodeTimeMs).average().orElse(0);
            System.out.printf("  平均编码时间: %.1f ms\n", avgKifEncode);
            System.out.printf("  平均解码时间: %.1f ms\n", avgKifDecode);
            
            // 与 PNG 对比
            List<ExperimentResult> pngResults = results.stream()
                .filter(r -> "png".equalsIgnoreCase(r.format))
                .collect(Collectors.toList());
            
            if (!pngResults.isEmpty() && kifResults.size() == pngResults.size()) {
                double pngAvg = pngResults.stream().mapToDouble(r -> r.compressionRatio).average().orElse(0);
                System.out.printf("  PNG 平均压缩率: %.2f%%\n", pngAvg);
                System.out.printf("  KIF vs PNG: %.2f%% (KIF 比 PNG %s %.2f%%)\n",
                    avgKifRatio,
                    avgKifRatio < pngAvg ? "小" : "大",
                    Math.abs(avgKifRatio - pngAvg));
                
                double pngAvgEncode = pngResults.stream().mapToLong(r -> r.encodeTimeMs).average().orElse(0);
                double pngAvgDecode = pngResults.stream().mapToLong(r -> r.decodeTimeMs).average().orElse(0);
                System.out.printf("  PNG 平均编码时间: %.1f ms\n", pngAvgEncode);
                System.out.printf("  PNG 平均解码时间: %.1f ms\n", pngAvgDecode);
                System.out.printf("  KIF 编码是 PNG 的 %.2f 倍\n", avgKifEncode / pngAvgEncode);
                System.out.printf("  KIF 解码是 PNG 的 %.2f 倍\n", avgKifDecode / pngAvgDecode);
            }
        }
    }

    /**
     * 实验结果记录类
     */
    static class ExperimentResult {
        String sourceFileName;
        String format;
        long originalSize;
        long compressedSize;
        double compressionRatio;
        long encodeTimeMs;
        long decodeTimeMs;
        boolean lossless;

        ExperimentResult(String sourceFileName, String format, long originalSize, 
                         long compressedSize, double compressionRatio, 
                         long encodeTimeMs, long decodeTimeMs, boolean lossless) {
            this.sourceFileName = sourceFileName;
            this.format = format;
            this.originalSize = originalSize;
            this.compressedSize = compressedSize;
            this.compressionRatio = compressionRatio;
            this.encodeTimeMs = encodeTimeMs;
            this.decodeTimeMs = decodeTimeMs;
            this.lossless = lossless;
        }

        @Override
        public String toString() {
            String status = lossless ? "✅ 无损" : "⚠️ 有损";
            return String.format("  %-6s | %,8d bytes (%.1f KB) | 压缩率: %5.1f%% | 编码: %4dms | 解码: %4dms | %s",
                format, compressedSize, compressedSize / 1024.0, 
                compressionRatio, encodeTimeMs, decodeTimeMs, status);
        }
    }
}