1"""
  2Cyclic Boundary Condition (CBC) Ring Generator for PyFock
  3This is a way to emulate the PBC
  4
  5Simple interface for creating ring supercells from 1D periodic systems and 
  6performing systematic convergence studies with extrapolation to thermodynamic limit.
  7
  8Usage:
  9    # Create ring supercells
 10    import PBC_ring
 11    ring_mol = PBC_ring.ring(unit_mol, N=10, periodicity=2.5)
 12    
 13    # Run DFT as usual
 14    basis = Basis(ring_mol, {'all': Basis.load(mol=ring_mol, basis_name='def2-SVP')})
 15    auxbasis = Basis(ring_mol, {'all': Basis.load(mol=ring_mol, basis_name='def2-universal-jfit')})
 16    dft_obj = DFT(ring_mol, basis, auxbasis, xc=[1, 7])
 17    energy = dft_obj.scf()
 18    
 19    # Convergence study utilities
 20    energies = PBC_ring.convergence_study(unit_mol, ring_sizes=[8,10,12,15], periodicity=2.5)
 21    tdl_energy = PBC_ring.extrapolate_tdl(energies)
 22"""
 23
 24import numpy as np
 25import math
 26import matplotlib.pyplot as plt
 27from typing import List, Tuple, Dict, Optional
 28from pyfock import Mol, Basis, DFT
 29
 30def ring(mol: Mol, N: int, periodicity: float, periodic_dir: str = 'x', 
 31         output_xyz: bool = True, xyz_filename: Optional[str] = None) -> Mol:
 32    """
 33    Create a ring supercell from a 1D periodic unit cell.
 34    
 35    Args:
 36        mol: PyFock Mol object containing the unit cell
 37        N: Number of unit cells to include in the ring
 38        periodicity: Length of unit cell along periodic direction (Angstrom)
 39        periodic_dir: Direction of periodicity ('x', 'y', or 'z'). Default: 'x'
 40        output_xyz: Whether to save XYZ file of the ring. Default: False
 41        xyz_filename: Custom filename for XYZ output. If None, auto-generates name
 42        
 43    Returns:
 44        Mol: New PyFock Mol object containing the ring structure
 45        
 46    Example:
 47        >>> unit_mol = Mol(coordfile='lih_unit.xyz')  # LiH unit cell
 48        >>> ring_mol = ring(unit_mol, N=10, periodicity=3.2)
 49        >>> # Now use ring_mol in normal PyFock DFT workflow
 50    """
 51    
 52    # Determine periodic direction index
 53    dir_map = {'x': 0, 'y': 1, 'z': 2}
 54    if periodic_dir.lower() not in dir_map:
 55        raise ValueError("periodic_dir must be 'x', 'y', or 'z'")
 56    periodic_idx = dir_map[periodic_dir.lower()]
 57    
 58    # Calculate ring radius from circumference = N * periodicity
 59    circumference = N * periodicity
 60    radius = circumference / (2.0 * math.pi)
 61    
 62    # Build ring coordinates
 63    ring_atoms = []
 64    
 65    for unit_idx in range(N):
 66        # Angular position for this unit
 67        unit_angle = 2.0 * math.pi * unit_idx / N
 68        
 69        # Add all atoms from this unit cell
 70        for atom in mol.atoms:
 71            symbol = atom[0]
 72            orig_coords = np.array([atom[1], atom[2], atom[3]])
 73            
 74            # Position along periodic direction within unit cell
 75            periodic_pos = orig_coords[periodic_idx]
 76            
 77            # Total angle for this atom
 78            atom_angle = unit_angle + (2.0 * math.pi * periodic_pos / periodicity) / N
 79            
 80            # Ring coordinates (place periodic direction in xy-plane)
 81            if periodic_idx == 0:  # x-direction periodic
 82                x_ring = radius * math.cos(atom_angle)
 83                y_ring = radius * math.sin(atom_angle)
 84                z_ring = orig_coords[2]
 85            elif periodic_idx == 1:  # y-direction periodic  
 86                x_ring = orig_coords[0]
 87                y_ring = radius * math.cos(atom_angle)
 88                z_ring = radius * math.sin(atom_angle)
 89            else:  # z-direction periodic
 90                x_ring = radius * math.cos(atom_angle)
 91                y_ring = orig_coords[1]  
 92                z_ring = radius * math.sin(atom_angle)
 93            
 94            ring_atoms.append([symbol, x_ring, y_ring, z_ring])
 95    
 96    # Create new Mol object for ring
 97    ring_mol = Mol(atoms=ring_atoms)
 98    ring_mol.label = f"Ring_{N}units_R{radius:.2f}A"
 99    
100    # Save XYZ file if requested
101    if output_xyz:
102        if xyz_filename is None:
103            xyz_filename = f"ring_{N}_units"
104        ring_mol.exportXYZ(xyz_filename, 
105                label=f"Ring with {N} units, R={radius:.3f} A, circumference={circumference:.3f} A")
106        print(f"Saved ring structure: {xyz_filename}")
107    
108    print(f"Created ring: {N} units, {len(ring_atoms)} atoms, radius={radius:.3f} A")
109    return ring_mol
110
111
112def ring_preserve_bonds(mol: Mol, N: int, periodicity: float, target_radius: float,
113                       periodic_dir: str = 'x', output_xyz: bool = True, 
114                       xyz_filename: Optional[str] = None) -> Mol:
115    """
116    Create a ring supercell that preserves interatomic spacing from the unit cell.
117    
118    This version prioritizes maintaining correct bond lengths over fitting exactly
119    around a complete circle. The ring may span less than or more than 2π radians.
120    
121    Args:
122        mol: PyFock Mol object containing the unit cell
123        N: Number of unit cells to include in the ring
124        periodicity: Length of unit cell along periodic direction (Angstrom)
125        target_radius: Desired radius for the ring (Angstrom)
126        periodic_dir: Direction of periodicity ('x', 'y', or 'z'). Default: 'x'
127        output_xyz: Whether to save XYZ file of the ring. Default: True
128        xyz_filename: Custom filename for XYZ output. If None, auto-generates name
129        
130    Returns:
131        Mol: New PyFock Mol object containing the ring structure
132        
133    Example:
134        >>> unit_mol = Mol(coordfile='lih_unit.xyz')  # LiH unit cell, 3.2 Å period
135        >>> # Create ring with 5 Å radius, preserving bond lengths
136        >>> ring_mol = ring_preserve_bonds(unit_mol, N=10, periodicity=3.2, target_radius=5.0)
137        >>> # Bonds remain 3.2 Å apart, but ring spans ~6.4 radians (not full 2π circle)
138    """
139    
140    # Determine periodic direction index
141    dir_map = {'x': 0, 'y': 1, 'z': 2}
142    if periodic_dir.lower() not in dir_map:
143        raise ValueError("periodic_dir must be 'x', 'y', or 'z'")
144    periodic_idx = dir_map[periodic_dir.lower()]
145    
146    # Calculate angular spacing to preserve bond lengths
147    angular_spacing_per_unit = periodicity / target_radius  # radians per unit cell
148    total_angle_spanned = N * angular_spacing_per_unit
149    
150    # Provide feedback about circle closure
151    angle_deficit = 2.0 * math.pi - total_angle_spanned
152    if abs(angle_deficit) > 0.1:  # More than ~6 degrees off
153        if angle_deficit > 0:
154            print(f"Note: Ring spans {total_angle_spanned:.3f} rad ({total_angle_spanned*180/math.pi:.1f}°)")
155            print(f"      Gap of {angle_deficit:.3f} rad ({angle_deficit*180/math.pi:.1f}°) to complete circle")
156        else:
157            print(f"Note: Ring spans {total_angle_spanned:.3f} rad ({total_angle_spanned*180/math.pi:.1f}°)")
158            print(f"      Overlaps by {-angle_deficit:.3f} rad ({-angle_deficit*180/math.pi:.1f}°) beyond full circle")
159    else:
160        print(f"Ring nearly closes: {total_angle_spanned:.3f} rad (~2π)")
161    
162    # Build ring coordinates
163    ring_atoms = []
164    
165    for unit_idx in range(N):
166        # Angular position for this unit (preserves spacing)
167        unit_angle = angular_spacing_per_unit * unit_idx
168        
169        # Add all atoms from this unit cell
170        for atom in mol.atoms:
171            symbol = atom[0]
172            orig_coords = np.array([atom[1], atom[2], atom[3]])
173            
174            # Position along periodic direction within unit cell
175            periodic_pos = orig_coords[periodic_idx]
176            
177            # Total angle for this atom (includes intra-unit-cell position)
178            atom_angle = unit_angle + (periodic_pos / target_radius)
179            
180            # Ring coordinates (place periodic direction in appropriate plane)
181            if periodic_idx == 0:  # x-direction periodic
182                x_ring = target_radius * math.cos(atom_angle)
183                y_ring = target_radius * math.sin(atom_angle)
184                z_ring = orig_coords[2]
185            elif periodic_idx == 1:  # y-direction periodic  
186                x_ring = orig_coords[0]
187                y_ring = target_radius * math.cos(atom_angle)
188                z_ring = target_radius * math.sin(atom_angle)
189            else:  # z-direction periodic
190                x_ring = target_radius * math.cos(atom_angle)
191                y_ring = orig_coords[1]  
192                z_ring = target_radius * math.sin(atom_angle)
193            
194            ring_atoms.append([symbol, x_ring, y_ring, z_ring])
195    
196    # Calculate actual circumference spanned
197    actual_circumference = total_angle_spanned * target_radius
198    
199    # Create new Mol object for ring
200    ring_mol = Mol(atoms=ring_atoms)
201    ring_mol.label = f"Ring_{N}units_R{target_radius:.2f}A_bondpreserved"
202    
203    # Save XYZ file if requested
204    if output_xyz:
205        if xyz_filename is None:
206            xyz_filename = f"ring_{N}units_R{target_radius:.1f}A_bonds_preserved.xyz"
207        
208        ring_mol.exportXYZ(xyz_filename, 
209                label=f"Bond-preserving ring: {N} units, R={target_radius:.3f} A, "
210                      f"span={total_angle_spanned:.3f} rad, arc_length={actual_circumference:.3f} A")
211        print(f"Saved ring structure: {xyz_filename}")
212    
213    print(f"Created bond-preserving ring: {N} units, {len(ring_atoms)} atoms")
214    print(f"  Radius: {target_radius:.3f} A")
215    print(f"  Arc length: {actual_circumference:.3f} A (vs {N*periodicity:.3f} A linear)")
216    print(f"  Bond spacing preserved: {periodicity:.3f} A")
217    
218    return ring_mol
219
220
221def suggest_radius_for_closure(N: int, periodicity: float) -> float:
222    """
223    Suggest a radius that would give near-perfect ring closure.
224    
225    Args:
226        N: Number of unit cells
227        periodicity: Unit cell length
228        
229    Returns:
230        float: Suggested radius for ~perfect circle closure
231    """
232    ideal_radius = (N * periodicity) / (2.0 * math.pi)
233    return ideal_radius
234
235
236def ring_with_closure_optimization(mol: Mol, N: int, periodicity: float, 
237                                 target_radius: Optional[float] = None,
238                                 optimize_closure: bool = True, 
239                                 periodic_dir: str = 'x',
240                                 output_xyz: bool = True,
241                                 xyz_filename: Optional[str] = None) -> Mol:
242    """
243    Create a ring with option to optimize for perfect closure or preserve bonds.
244    
245    Args:
246        mol: PyFock Mol object containing the unit cell
247        N: Number of unit cells
248        periodicity: Unit cell length (Angstrom)
249        target_radius: Desired radius. If None and optimize_closure=True, calculates optimal
250        optimize_closure: If True and target_radius=None, optimizes for perfect ring closure
251        periodic_dir: Direction of periodicity ('x', 'y', or 'z')
252        output_xyz: Whether to save XYZ file
253        xyz_filename: Custom filename for XYZ output
254        
255    Returns:
256        Mol: Ring structure
257    """
258    
259    if target_radius is None:
260        if optimize_closure:
261            target_radius = suggest_radius_for_closure(N, periodicity)
262            print(f"Using closure-optimized radius: {target_radius:.3f} A")
263        else:
264            raise ValueError("Must provide target_radius if optimize_closure=False")
265    
266    return ring_preserve_bonds(mol, N, periodicity, target_radius, 
267                             periodic_dir, output_xyz, xyz_filename)