# Author: Lukas Maier
# Affiliation: Institute of Process and Particle Engineering, TU Graz
# License: LGPL 3.0

"""
Module for extracting vertical bounds (min/max Z) from LIGGGHTS simulation data.

This script parses LIGGGHTS dump files to find the highest and lowest points occupied
by particles, accounting for their radii. This information is typically used to
adjust simulation boundaries or define compaction limits.
"""

import os
import glob
import math
from collections import defaultdict

def extract_max_min_z_with_radius_adjusted(folder_path, debug=False):
    """
    Extracts the maximum and minimum Z coordinates from particle data in a given folder.

    Scans the folder for the latest LIGGGHTS dump files matching 'spheres_*' or
    'multisphere_*', reads particle positions and radii, and calculates the absolute
    Z bounds of the particle assembly.

    Args:
        folder_path (str): Path to the directory containing LIGGGHTS dump files.
        debug (bool): If True, prints detailed debugging information.

    Returns:
        tuple: A tuple containing (max_z, min_z). Respective values are floats representing
              the highest and lowest Z-coordinates adjusted for particle radius.
              Returns (None, None) if no valid data is found.
    """
    coordinates = []
    file_patterns = ["spheres_*", "multisphere_*"]
    
    if debug:
        print(f"DEBUG: Looking for files in folder: {folder_path}")
        print(f"DEBUG: File patterns: {file_patterns}")
    
    # Read coordinates and radii from files
    for pattern in file_patterns:
        matching_files = glob.glob(os.path.join(folder_path, pattern))
        if debug:
            print(f"DEBUG: Pattern '{pattern}' found {len(matching_files)} files: {matching_files}")
        
        if matching_files:
            latest_file = max(matching_files, key=os.path.getmtime)
            if debug:
                print(f"DEBUG: Using latest file: {latest_file}")
            
            with open(latest_file, 'r') as file:
                lines = file.readlines()
                if debug:
                    print(f"DEBUG: File {latest_file} has {len(lines)} lines")
                    print(f"DEBUG: First few lines: {lines[:5]}")
                
                for line_num, line in enumerate(lines):
                    if line.strip() and not line.startswith("ITEM:"):
                        parts = line.split()
                        if debug and line_num < 5:
                            print(f"DEBUG: Line {line_num}: parts={parts}")
                        try:
                            x = float(parts[6])
                            y = float(parts[7])
                            z = float(parts[8])
                            radius = float(parts[3])
                            coordinates.append((x, y, z, radius))
                        except (IndexError, ValueError) as e:
                            if debug and line_num < 5:
                                print(f"DEBUG: Failed to parse line {line_num}: {e}")
                            continue
    
    if debug:
        print(f"DEBUG: Total coordinates extracted: {len(coordinates)}")
        if coordinates:
            print(f"DEBUG: First few coordinates: {coordinates[:5]}")
        else:
            print("DEBUG: No coordinates found!")
            return None, None

    # Use a spatial grid to group atoms into clusters (exclude single particles)
    # Use a multiple of particle radius for neighbor detection
    radius_multiplier = 3.0  # More permissive: allow particles to be further apart to form clusters
    
    # Calculate average radius for grid sizing
    avg_radius = sum(radius for _, _, _, radius in coordinates) / len(coordinates) if coordinates else 0.001
    grid_size = avg_radius * radius_multiplier
    grid = defaultdict(list)

    if debug:
        print(f"DEBUG: Average radius: {avg_radius}")
        print(f"DEBUG: Grid size: {grid_size}")
        print(f"DEBUG: Radius multiplier: {radius_multiplier}")

    for i, (x, y, z, radius) in enumerate(coordinates):
        grid_key = (int(x // grid_size), int(y // grid_size), int(z // grid_size))
        grid[grid_key].append((x, y, z, radius, i))

    if debug:
        print(f"DEBUG: Created {len(grid)} grid cells")
        print(f"DEBUG: Grid cells with particles: {[(k, len(v)) for k, v in list(grid.items())[:5]]}")

    clusters = []
    visited = set()

    def find_neighbors(x1, y1, z1, radius1, idx1):
        """Find all neighbors of a given particle using BFS"""
        neighbors = []
        neighbor_count = 0
        # Check neighboring grid cells
        grid_key = (int(x1 // grid_size), int(y1 // grid_size), int(z1 // grid_size))
        for dx in [-1, 0, 1]:
            for dy in [-1, 0, 1]:
                for dz in [-1, 0, 1]:
                    neighbor_key = (grid_key[0] + dx, grid_key[1] + dy, grid_key[2] + dz)
                    for x2, y2, z2, radius2, idx2 in grid.get(neighbor_key, []):
                        if idx2 != idx1 and idx2 not in visited:
                            distance = math.sqrt((x1 - x2)**2 + (y1 - y2)**2 + (z1 - z2)**2)
                            # Use radius-based threshold: particles are neighbors if they're close enough
                            threshold = radius_multiplier * (radius1 + radius2) / 2
                            if distance <= threshold:
                                neighbors.append((x2, y2, z2, radius2, idx2))
                                neighbor_count += 1
        
        if debug and idx1 < 5:  # Debug first few particles
            print(f"DEBUG: Particle {idx1} at ({x1:.6f}, {y1:.6f}, {z1:.6f}) found {neighbor_count} neighbors")
            if neighbor_count > 0:
                distances = []
                for x2, y2, z2, radius2, idx2 in neighbors[:3]:  # Show first 3 neighbors
                    dist = math.sqrt((x1 - x2)**2 + (y1 - y2)**2 + (z1 - z2)**2)
                    distances.append(dist)
                print(f"DEBUG: Sample neighbor distances: {distances}")
                print(f"DEBUG: Threshold used: {radius_multiplier * (radius1 + avg_radius) / 2:.6f}")
        
        return neighbors

    for grid_key, atoms in grid.items():
        for x1, y1, z1, radius1, idx1 in atoms:
            if idx1 in visited:
                continue
            
            # Use BFS to find all connected particles
            cluster = [(x1, y1, z1, radius1)]
            visited.add(idx1)
            queue = [(x1, y1, z1, radius1, idx1)]
            
            while queue:
                curr_x, curr_y, curr_z, curr_radius, curr_idx = queue.pop(0)
                neighbors = find_neighbors(curr_x, curr_y, curr_z, curr_radius, curr_idx)
                
                for x2, y2, z2, radius2, idx2 in neighbors:
                    if idx2 not in visited:
                        cluster.append((x2, y2, z2, radius2))
                        visited.add(idx2)
                        queue.append((x2, y2, z2, radius2, idx2))
            
            clusters.append(cluster)

    if debug:
        print(f"DEBUG: Found {len(clusters)} initial clusters")
        cluster_sizes = [len(cluster) for cluster in clusters]
        if cluster_sizes:
            print(f"DEBUG: Cluster sizes: {sorted(cluster_sizes, reverse=True)[:10]}")

    # Filter clusters to exclude outliers more effectively
    min_cluster_size = 3  # Require at least 3 particles to form a valid cluster
    large_clusters = [cluster for cluster in clusters if len(cluster) >= min_cluster_size]
    
    if debug:
        print(f"DEBUG: After min size filter ({min_cluster_size}): {len(large_clusters)} clusters")
    
    # Additional filter: only consider clusters that are at least 10% the size of the largest cluster
    if large_clusters:
        largest_cluster_size = max(len(cluster) for cluster in large_clusters)
        size_threshold = max(min_cluster_size, int(0.1 * largest_cluster_size))
        multi_particle_clusters = [cluster for cluster in large_clusters if len(cluster) >= size_threshold]
        
        if debug:
            print(f"DEBUG: Largest cluster size: {largest_cluster_size}")
            print(f"DEBUG: Size threshold (10%): {size_threshold}")
            print(f"DEBUG: After size threshold filter: {len(multi_particle_clusters)} clusters")
    else:
        multi_particle_clusters = large_clusters
    
    if debug:
        print(f"DEBUG: Multi-particle clusters before outlier removal: {len(multi_particle_clusters)}")
    
    # Statistical outlier removal: remove particles with z-coordinates that are statistical outliers
    if multi_particle_clusters:
        # Collect all z-coordinates from valid clusters
        all_z_coords = []
        for cluster in multi_particle_clusters:
            for x, y, z, radius in cluster:
                all_z_coords.append(z)
        
        if debug:
            print(f"DEBUG: Collected {len(all_z_coords)} z-coordinates")
            if all_z_coords:
                print(f"DEBUG: Z-coord range: {min(all_z_coords):.3f} to {max(all_z_coords):.3f}")
        
        if all_z_coords:
            # Calculate statistics
            import statistics
            mean_z = statistics.mean(all_z_coords)
            stdev_z = statistics.stdev(all_z_coords) if len(all_z_coords) > 1 else 0
            
            if debug:
                print(f"DEBUG: Z statistics - mean: {mean_z:.3f}, stdev: {stdev_z:.3f}")
            
            # Remove outliers beyond 2 standard deviations
            z_threshold = 2.8
            z_min_bound = mean_z - z_threshold * stdev_z
            z_max_bound = mean_z + z_threshold * stdev_z
            
            if debug:
                print(f"DEBUG: Z bounds for outlier removal: {z_min_bound:.3f} to {z_max_bound:.3f}")
            
            # Filter out particles that are statistical outliers
            filtered_clusters = []
            outlier_count = 0
            for cluster in multi_particle_clusters:
                filtered_cluster = []
                for x, y, z, radius in cluster:
                    if z_min_bound <= z <= z_max_bound:
                        filtered_cluster.append((x, y, z, radius))
                    else:
                        outlier_count += 1
                
                # Only keep clusters that still have particles after outlier removal
                if len(filtered_cluster) >= min_cluster_size:
                    filtered_clusters.append(filtered_cluster)
            
            if debug:
                print(f"DEBUG: Removed {outlier_count} outlier particles")
                print(f"DEBUG: Final filtered clusters: {len(filtered_clusters)}")
            
            multi_particle_clusters = filtered_clusters

    # Calculate max and min z values from all filtered multi-particle clusters
    max_z = float('-inf')
    min_z = float('inf')
    
    if debug:
        print(f"DEBUG: Calculating max/min Z from {len(multi_particle_clusters)} final clusters")
    
    total_particles_used = 0
    for cluster in multi_particle_clusters:
        for x, y, z, radius in cluster:
            total_particles_used += 1
            if z + radius > max_z:
                max_z = z + radius
            if z - radius < min_z:
                min_z = z - radius
    
    if debug:
        print(f"DEBUG: Used {total_particles_used} particles for max/min calculation")
        print(f"DEBUG: Raw max_z: {max_z}, Raw min_z: {min_z}")
    
    if max_z == float('-inf') or min_z == float('inf'):
        if debug:
            print("DEBUG: No valid particles found - returning None, None")
        return None, None
    
    if debug:
        print(f"DEBUG: Final results - MAX_Z: {max_z}, MIN_Z: {min_z}")
    
    return max_z, min_z

if __name__ == "__main__":
    post_fill_folder = os.path.join("post_states")
    max_z_value, min_z_value = extract_max_min_z_with_radius_adjusted(post_fill_folder,debug=False)
    # Output only the variables for shell evaluation
    print(f"MAX_Z={max_z_value}")
    print(f"MIN_Z={min_z_value}")
