#!/usr/bin/env python3
"""
Program to analyze JSON files produced by rule_parser.py and output statistics about:
1. Number of loop formulas
2. Number of atoms with underlines (_) in their names
3. Number of rules where atoms with underlines appear as heads
4. Number of rules where atoms with underlines appear in bodies
"""

import json
import os
import sys
import argparse
from pathlib import Path


def analyze_json_file(json_path):
    """
    Analyze a single JSON file and return statistics.

    Args:
        json_path: Path to the JSON file

    Returns:
        Dictionary with analysis results
    """
    try:
        with open(json_path, 'r') as f:
            data = json.load(f)
    except (FileNotFoundError, json.JSONDecodeError) as e:
        print(f"Error reading {json_path}: {e}")
        return None

    # Count loop formulas
    loop_formulas_count = len(data.get('loop_formulas', []))

    # Find atoms with underlines in their names
    atom_mapping = data.get('atom_mapping', {})
    auxiliary_atoms = {}  # atom_id -> atom_name

    for atom_id, atom_name in atom_mapping.items():
        if '_' in atom_name:
            auxiliary_atoms[int(atom_id)] = atom_name

    auxiliary_atoms_count = len(auxiliary_atoms)

    # Count rules where auxiliary atoms appear as heads
    normal_rules = data.get('rules', {}).get('normal', [])
    disjunctive_rules = data.get('rules', {}).get('disjunctive', [])
    choice_rules = data.get('rules', {}).get('choice', [])

    aux_head_rules_count = 0
    aux_body_rules_count = 0

    # Analyze normal rules
    for rule in normal_rules:
        head = rule.get('head')
        body = rule.get('body', {})
        pos_body = body.get('pos', [])
        neg_body = body.get('neg', [])

        # Check if head is an auxiliary atom
        if head in auxiliary_atoms:
            aux_head_rules_count += 1

        # Check if any atom in body is auxiliary
        all_body_atoms = pos_body + neg_body
        if any(atom_id in auxiliary_atoms for atom_id in all_body_atoms):
            aux_body_rules_count += 1

    # Analyze disjunctive rules
    for rule in disjunctive_rules:
        head = rule.get('head', [])
        body = rule.get('body', {})
        pos_body = body.get('pos', [])
        neg_body = body.get('neg', [])

        # Check if any head atom is auxiliary
        if any(atom_id in auxiliary_atoms for atom_id in head):
            aux_head_rules_count += 1

        # Check if any atom in body is auxiliary
        all_body_atoms = pos_body + neg_body
        if any(atom_id in auxiliary_atoms for atom_id in all_body_atoms):
            aux_body_rules_count += 1

    # Analyze choice rules
    for rule in choice_rules:
        head = rule.get('head', [])
        body = rule.get('body', {})
        pos_body = body.get('pos', [])
        neg_body = body.get('neg', [])

        # Check if any head atom is auxiliary
        if any(atom_id in auxiliary_atoms for atom_id in head):
            aux_head_rules_count += 1

        # Check if any atom in body is auxiliary
        all_body_atoms = pos_body + neg_body
        if any(atom_id in auxiliary_atoms for atom_id in all_body_atoms):
            aux_body_rules_count += 1

    return {
        'loop_formulas_count': loop_formulas_count,
        'auxiliary_atoms_count': auxiliary_atoms_count,
        'aux_head_rules_count': aux_head_rules_count,
        'aux_body_rules_count': aux_body_rules_count,
        'auxiliary_atoms': auxiliary_atoms
    }


def analyze_directory(directory_path, verbose=False):
    """
    Analyze all subdirectories in the given directory path.
    For each subdirectory, look for a JSON file with the same name.

    Args:
        directory_path: Path to the directory containing subdirectories to analyze
        verbose: Whether to print detailed information
    """
    directory_path = Path(directory_path)

    if not directory_path.exists() or not directory_path.is_dir():
        print(f"Error: {directory_path} is not a valid directory")
        return

    results = []

    # Iterate through all subdirectories
    for subdir in sorted(directory_path.iterdir()):
        # if the name starts with pin
        if not subdir.name.startswith('pin'):
            continue

        if subdir.is_dir():
            # Look for JSON file with the same name as the subdirectory
            json_file = subdir / f"{subdir.name}.json"

            if json_file.exists():
                if verbose:
                    print(f"Analyzing {json_file}...")

                analysis = analyze_json_file(json_file)
                if analysis:
                    results.append({
                        'directory': subdir.name,
                        'json_file': str(json_file),
                        **analysis
                    })
            else:
                if verbose:
                    print(f"No JSON file found for {subdir.name}")

    return results


def print_results(results, show_atoms=False):
    """
    Print the analysis results in a formatted table.

    Args:
        results: List of analysis results
        show_atoms: Whether to show the auxiliary atoms details
    """
    if not results:
        print("No results to display.")
        return

    # Print header
    print(f"{'Directory':<20} {'Loop Formulas':<15} {'Aux Atoms':<12} {'Aux Head Rules':<15} {'Aux Body Rules':<15}")
    print("-" * 80)

    # Print results
    for result in results:
        print(f"{result['directory']:<20} "
              f"{result['loop_formulas_count']:<15} "
              f"{result['auxiliary_atoms_count']:<12} "
              f"{result['aux_head_rules_count']:<15} "
              f"{result['aux_body_rules_count']:<15}")

        if show_atoms and result['auxiliary_atoms']:
            print(f"  Auxiliary atoms: {list(result['auxiliary_atoms'].values())}")

    # Print summary
    total_loops = sum(r['loop_formulas_count'] for r in results)
    total_aux_atoms = sum(r['auxiliary_atoms_count'] for r in results)
    total_aux_head_rules = sum(r['aux_head_rules_count'] for r in results)
    total_aux_body_rules = sum(r['aux_body_rules_count'] for r in results)

    print("-" * 80)
    print(f"{'TOTAL':<20} "
          f"{total_loops:<15} "
          f"{total_aux_atoms:<12} "
          f"{total_aux_head_rules:<15} "
          f"{total_aux_body_rules:<15}")


def main():
    parser = argparse.ArgumentParser(
        description="Analyze JSON files produced by rule_parser.py for loop formulas and auxiliary atoms statistics"
    )
    parser.add_argument(
        "directory",
        help="Directory containing subdirectories with JSON files to analyze"
    )
    parser.add_argument(
        "-v", "--verbose",
        action="store_true",
        help="Enable verbose output"
    )
    parser.add_argument(
        "-a", "--show-atoms",
        action="store_true",
        help="Show detailed auxiliary atoms information"
    )
    parser.add_argument(
        "-f", "--file",
        help="Analyze a single JSON file instead of a directory"
    )

    args = parser.parse_args()

    if args.file:
        # Analyze single file
        analysis = analyze_json_file(args.file)
        if analysis:
            print(f"Analysis of {args.file}:")
            print(f"Loop formulas: {analysis['loop_formulas_count']}")
            print(f"Auxiliary atoms: {analysis['auxiliary_atoms_count']}")
            print(f"Rules with auxiliary atoms as heads: {analysis['aux_head_rules_count']}")
            print(f"Rules with auxiliary atoms in bodies: {analysis['aux_body_rules_count']}")

            if args.show_atoms and analysis['auxiliary_atoms']:
                print("\nAuxiliary atoms:")
                for atom_id, atom_name in analysis['auxiliary_atoms'].items():
                    print(f"  {atom_id}: {atom_name}")
    else:
        # Analyze directory
        results = analyze_directory(args.directory, args.verbose)
        if results:
            print_results(results, args.show_atoms)
        else:
            print("No JSON files found to analyze.")


if __name__ == "__main__":
    main()
