#!/usr/bin/env python3
"""
JSON to CSV Converter for VeriCluX Data

This script converts JSON files in the data directory to CSV format for easier analysis.
It handles different types of explanation data (individual, cluster, inflated, robust).
"""

import json
from pathlib import Path
import pandas as pd
import argparse


class JSONToCSVConverter:
    def __init__(self, data_dir: str = "data"):
        self.data_dir = Path(data_dir)

    def convert_individual_explanations(self, input_dir: Path, output_file: Path):
        """Convert individual explanations to CSV"""
        data = []

        for json_file in input_dir.glob("*.json"):
            with open(json_file, 'r') as f:
                json_data = json.load(f)

            # Extract basic metadata
            row = {
                'dataset': json_data.get('dataset', ''),
                'custom_dim': json_data.get('custom_dim', 0),
                'label': json_data.get('label', ''),
                'epsilon': json_data.get('epsilon', ''),
                'ordering': json_data.get('ordering', ''),
                'seed': json_data.get('seed', ''),
                'id': json_data.get('id', ''),
                'sample_index': json_data.get('sample_index', ''),
                'computation_time': json_data.get('computation_time', ''),
                'sample': str(json_data.get('sample', [])),
                'explanations': str(json_data.get('explanations', []))
            }
            data.append(row)

        if data:
            df = pd.DataFrame(data)
            df.to_csv(output_file, index=False)
            print(f"Converted {len(data)} individual explanation files to {output_file}")

    def convert_cluster_explanations(self, input_dir: Path, output_file: Path):
        """Convert cluster explanations to CSV"""
        data = []

        for json_file in input_dir.glob("*.json"):
            with open(json_file, 'r') as f:
                json_data = json.load(f)

            # Extract basic metadata
            row = {
                'dataset': json_data.get('dataset', ''),
                'custom_dim': json_data.get('custom_dim', 0),
                'label': json_data.get('label', ''),
                'epsilon': json_data.get('epsilon', ''),
                'ordering': json_data.get('ordering', ''),
                'seed': json_data.get('seed', ''),
                'cluster_index': json_data.get('cluster_index', ''),
                'cluster_size': json_data.get('cluster_size', ''),
                'id': json_data.get('id', ''),
                'computation_time': json_data.get('computation_time', ''),
                'delta': json_data.get('delta', ''),
                'prototype': json_data.get('prototype', ''),
                'samples': str(json_data.get('samples', [])),
                'explanations': str(json_data.get('explanations', []))
            }
            data.append(row)

        if data:
            df = pd.DataFrame(data)
            df.to_csv(output_file, index=False)
            print(f"Converted {len(data)} cluster explanation files to {output_file}")

    def convert_inflated_explanations(self, input_dir: Path, output_file: Path):
        """Convert inflated explanations to CSV"""
        data = []

        for json_file in input_dir.glob("*.json"):
            with open(json_file, 'r') as f:
                json_data = json.load(f)

            # Extract basic metadata
            row = {
                'dataset': json_data.get('dataset', ''),
                'custom_dim': json_data.get('custom_dim', 0),
                'label': json_data.get('label', ''),
                'epsilon': json_data.get('epsilon', ''),
                'ordering': json_data.get('ordering', ''),
                'seed': json_data.get('seed', ''),
                'sample_index': json_data.get('sample_index', ''),
                'id': json_data.get('id', ''),
                'computation_time': json_data.get('computation_time', ''),
                'delta': json_data.get('delta', ''),
                'prototype': json_data.get('prototype', ''),
                'eta': json_data.get('eta', ''),
                'samples': str(json_data.get('samples', [])),
                'explanations': str(json_data.get('explanations', [])),
                'inflated_explanations': str(json_data.get('inflated_explanations', []))
            }
            data.append(row)

        if data:
            df = pd.DataFrame(data)
            df.to_csv(output_file, index=False)
            print(f"Converted {len(data)} inflated explanation files to {output_file}")

    def convert_robust_explanations(self, input_dir: Path, output_file: Path):
        """Convert robust explanations to CSV"""
        data = []

        for json_file in input_dir.glob("*.json"):
            with open(json_file, 'r') as f:
                json_data = json.load(f)

            # Extract basic metadata
            row = {
                'dataset': json_data.get('dataset', ''),
                'custom_dim': json_data.get('custom_dim', 0),
                'label': json_data.get('label', ''),
                'epsilon': json_data.get('epsilon', ''),
                'ordering': json_data.get('ordering', ''),
                'seed': json_data.get('seed', ''),
                'sample_index': json_data.get('sample_index', ''),
                'id': json_data.get('id', ''),
                'computation_time': json_data.get('computation_time', ''),
                'delta': json_data.get('delta', ''),
                'prototype': json_data.get('prototype', ''),
                'samples': str(json_data.get('samples', [])),
                'explanations': str(json_data.get('explanations', []))
            }
            data.append(row)

        if data:
            df = pd.DataFrame(data)
            df.to_csv(output_file, index=False)
            print(f"Converted {len(data)} robust explanation files to {output_file}")

    def convert_all(self, output_dir: str = "data/explanations"):
        """Convert all JSON files in the data directory to CSV"""
        output_path = Path(output_dir)
        output_path.mkdir(exist_ok=True, parents=True)

        # Process delta_robust explanations
        delta_robust_dir = self.data_dir / "delta_robust"
        if delta_robust_dir.exists():
            # Classic datasets
            classic_dir = delta_robust_dir / "classic"
            if classic_dir.exists():
                classic_output = output_path / "delta_robust" / "classic"
                classic_output.mkdir(exist_ok=True, parents=True)

                # Inflated explanations
                inflated_dir = classic_dir / "inflated-explanations"
                if inflated_dir.exists():
                    self.convert_inflated_explanations(
                        inflated_dir,
                        classic_output / "inflated_explanations.csv"
                    )

                # Robust explanations
                robust_dir = classic_dir / "robust-explanations"
                if robust_dir.exists():
                    self.convert_robust_explanations(
                        robust_dir,
                        classic_output / "robust_explanations.csv"
                    )

            # Scalability datasets
            scalability_dir = delta_robust_dir / "scalability"
            if scalability_dir.exists():
                scalability_output = output_path / "delta_robust" / "scalability"
                scalability_output.mkdir(exist_ok=True, parents=True)

                # Inflated explanations
                inflated_dir = scalability_dir / "inflated-explanations"
                if inflated_dir.exists():
                    self.convert_inflated_explanations(
                        inflated_dir,
                        scalability_output / "inflated_explanations.csv"
                    )

                # Robust explanations
                robust_dir = scalability_dir / "robust-explanations"
                if robust_dir.exists():
                    self.convert_robust_explanations(
                        robust_dir,
                        scalability_output / "robust_explanations.csv"
                    )

        # Process cluster explanations
        cluster_dir = self.data_dir / "cluster"
        if cluster_dir.exists():
            # Classic datasets
            classic_dir = cluster_dir / "classic"
            if classic_dir.exists():
                classic_output = output_path / "cluster" / "classic"
                classic_output.mkdir(exist_ok=True, parents=True)

                # Individual explanations
                individual_dir = classic_dir / "individual-explanations"
                if individual_dir.exists():
                    self.convert_individual_explanations(
                        individual_dir,
                        classic_output / "individual_explanations.csv"
                    )

                # Cluster explanations
                cluster_explanations_dir = classic_dir / "cluster-explanations"
                if cluster_explanations_dir.exists():
                    self.convert_cluster_explanations(
                        cluster_explanations_dir,
                        classic_output / "cluster_explanations.csv"
                    )

            # Scalability datasets
            scalability_dir = cluster_dir / "scalability"
            if scalability_dir.exists():
                scalability_output = output_path / "cluster" / "scalability"
                scalability_output.mkdir(exist_ok=True, parents=True)

                # Individual explanations
                individual_dir = scalability_dir / "individual-explanations"
                if individual_dir.exists():
                    self.convert_individual_explanations(
                        individual_dir,
                        scalability_output / "individual_explanations.csv"
                    )

                # Cluster explanations by Number of Features
                cluster_explanations_dir = scalability_dir / "cluster-explanations-by-num-features"
                if cluster_explanations_dir.exists():
                    self.convert_cluster_explanations(
                        cluster_explanations_dir,
                        scalability_output / "cluster_explanations_by_num_features.csv"
                    )

                # Cluster explanations by Group Size
                cluster_explanations_dir = scalability_dir / "cluster-explanations-by-cluster-size"
                if cluster_explanations_dir.exists():
                    self.convert_cluster_explanations(
                        cluster_explanations_dir,
                        scalability_output / "cluster_explanations_by_cluster_size.csv"
                    )

            # Examples
            examples_dir = cluster_dir / "examples"
            if examples_dir.exists():
                examples_output = output_path / "cluster" / "examples"
                examples_output.mkdir(exist_ok=True, parents=True)

                # Individual explanations
                individual_dir = examples_dir / "individual-explanations"
                if individual_dir.exists():
                    self.convert_individual_explanations(
                        individual_dir,
                        examples_output / "individual_explanations.csv"
                    )

                # Cluster explanations
                cluster_explanations_dir = examples_dir / "cluster-explanations"
                if cluster_explanations_dir.exists():
                    self.convert_cluster_explanations(
                        cluster_explanations_dir,
                        examples_output / "cluster_explanations.csv"
                    )

                # Ordering explanations
                ordering_explanations_dir = examples_dir / "ordering-explanations"
                if ordering_explanations_dir.exists():
                    self.convert_individual_explanations(
                        ordering_explanations_dir,
                        examples_output / "ordering_explanations.csv"
                    )

        print(f"\nAll conversions completed! CSV files saved to {output_path}")


def main():
    args = parse_arguments()

    converter = JSONToCSVConverter(args.data_dir)
    converter.convert_all(args.output_dir)


def parse_arguments():
    parser = argparse.ArgumentParser(description="Convert VeriCluX JSON data to CSV format")
    parser.add_argument("--data-dir", default="data/logs", help="Input data directory (default: data/logs)")
    parser.add_argument("--output-dir", default="data/explanations", help="Output CSV directory (default: data/explanations)")
    return parser.parse_args()


if __name__ == "__main__":
    main()
