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

"""
Module for coordinate extraction and boundary analysis of LIGGGHTS simulations.

This script analyzes particle positions from simulation dump files to determine
the maximum spatial extent (maximum absolute coordinate) of the system. It handles
particle clustering and checks for overlaps with defined boundaries.
"""

import os
import glob
import math
from collections import defaultdict

def extract_max_absolute_coordinate_with_radius(folder_path, debug=False):
    """
    Calculates the maximum absolute coordinate value across all particles.

    Reads the latest dump files for spheres and multispheres, parses particle
    positions and radii, and builds clusters to find the spatial extent of
    the system.

    Args:
        folder_path (str): Path to the directory containing LIGGGHTS dump files.
        debug (bool): If True, enables debug logging.

    Returns:
        float: The maximum absolute coordinate value (considering x, y, z and radius)
               found among all particles. Returns 0 if no coordinates are found.
    """
    coordinates = []
    file_patterns = ["spheres_*", "multisphere_*"]
    
    # Read coordinates and radii from files
    for pattern in file_patterns:
        for file_path in glob.glob(os.path.join(folder_path, pattern)):
            latest_file = max(glob.glob(os.path.join(folder_path, pattern)), key=os.path.getmtime)
            file_path = latest_file
            if debug:
                print(f"Processing file: {file_path}")
            with open(file_path, 'r') as file:
                lines = file.readlines()
                for line in lines:
                    if line.strip() and not line.startswith("ITEM:"):
                        parts = line.split()
                        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):
                            if debug:
                                print(f"Skipping line due to error: {line.strip()}")
                            continue

    # Use a spatial grid to group atoms into clusters
    threshold_distance = 0.0005  # Define a fixed threshold distance
    grid_size = threshold_distance
    grid = defaultdict(list)

    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))

    clusters = []
    visited = set()

    for grid_key, atoms in grid.items():
        for x1, y1, z1, radius1, idx1 in atoms:
            if idx1 in visited:
                continue
            cluster = [(x1, y1, z1, radius1)]
            visited.add(idx1)
            # Check neighboring grid cells
            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 not in visited:
                                distance = math.sqrt((x1 - x2)**2 + (y1 - y2)**2 + (z1 - z2)**2)
                                if distance <= threshold_distance:
                                    cluster.append((x2, y2, z2, radius2))
                                    visited.add(idx2)
            clusters.append(cluster)

    # Find the largest cluster
    largest_cluster = max(clusters, key=len, default=[])

    # Calculate the maximum absolute coordinate from the largest cluster
    if largest_cluster:
        max_abs_coord = max(
            max(abs(x) + radius, abs(y) + radius, abs(z) + radius)
            for x, y, z, radius in largest_cluster
        )
    else:
        if debug:
            print("No valid clusters found.")
        max_abs_coord = float('-inf')

    return max_abs_coord

def check_boundary_overlap(folder_path, boundary_value, debug=False):
    """
    Checks if any particle extends beyond a specified spatial boundary.

    Iterates through particle data to determine if any particle's position plus
    its radius exceeds the given boundary value in any dimension (X, Y, or Z).

    Args:
        folder_path (str): Path to the directory containing LIGGGHTS dump files.
        boundary_value (float): The coordinate limit to check against.
        debug (bool): If True, enables debug logging regarding overlaps.

    Returns:
        bool: True if any particle overlaps with or exceeds the boundary, False otherwise.
    """
    file_patterns = ["spheres_*", "multisphere_*"]
    
    for pattern in file_patterns:
        for file_path in glob.glob(os.path.join(folder_path, pattern)):
            latest_file = max(glob.glob(os.path.join(folder_path, pattern)), key=os.path.getmtime)
            file_path = latest_file
            with open(file_path, 'r') as file:
                lines = file.readlines()
                for line in lines:
                    if line.strip() and not line.startswith("ITEM:"):
                        parts = line.split()
                        try:
                            x = float(parts[6])
                            y = float(parts[7])
                            z = float(parts[8])
                            radius = float(parts[3])
                            
                            # Check if particle overlaps with any boundary
                            if (abs(x) + radius > boundary_value or 
                                abs(y) + radius > boundary_value or 
                                abs(z) + radius > boundary_value):
                                if debug:
                                    print(f"Particle at ({x}, {y}, {z}) with radius {radius} overlaps with boundary {boundary_value}")
                                return True
                        except (IndexError, ValueError):
                            continue
    
    return False

if __name__ == "__main__":
    post_fill_folder = os.path.join("post_states")
    max_abs_coord_value = extract_max_absolute_coordinate_with_radius(post_fill_folder)
    
    # Initial scaling factor
    scaling_factor = 1.05
    max_abs_coord_value *= scaling_factor
    
    # Check for boundary overlaps and incrementally increase boundary if needed
    has_overlap = check_boundary_overlap(post_fill_folder, max_abs_coord_value)
    
    # Parameters for boundary enlargement
    max_attempts = 50
    increment = 0.01  # Increase by 1% each time
    attempt = 0
    
    while has_overlap and attempt < max_attempts:
        #print(f"WARNING: Some particles overlap with the boundary value of {max_abs_coord_value}")
        scaling_factor += increment
        max_abs_coord_value = extract_max_absolute_coordinate_with_radius(post_fill_folder) * scaling_factor
        #print(f"Enlarging boundary by {increment:.2%} to {max_abs_coord_value} (scaling factor: {scaling_factor:.2f})")
        has_overlap = check_boundary_overlap(post_fill_folder, max_abs_coord_value)
        attempt += 1
    
    # Apply final safety margin after resolving overlaps
    print(f"MAX_ABS_COORD={max_abs_coord_value}")
