# compute_density.tcl
# Computes the overall mass density (g/cm³) per frame from a NAMD trajectory.
# Requires a PSF and DCD loaded in VMD with periodic cell info.
#
# Author:
# Diego E. B. Gomes
#
# Usage:
#   source compute_density.tcl
#   compute_density top "density_output.dat"

proc compute_density {molid outfile} {
    set nframes [molinfo $molid get numframes]
    set sel [atomselect $molid "all"]

    # Compute total mass once (doesn't change across frames)
    set total_mass [vecsum [$sel get mass]]

    # Convert: mass in amu, volume in ų
    # 1 amu = 1.66053906660e-24 g
    # 1 ų  = 1e-24 cm³
    # density = (total_mass * 1.66053906660e-24) / (volume * 1e-24)
    #         = total_mass * 1.66053906660 / volume
    set amu_to_g 1.66053906660

    set pi 3.14159265358979
    
    # set pi/180 to avoid repeated computation
    set pi180  [expr $pi / 180.0 ]   

    set fp [open $outfile w]
    puts $fp "# frame  density(g/cm3)  volume(A3)  a  b  c"

    set sum 0.0

    for {set i 0} {$i < $nframes} {incr i} {
        $sel frame $i
        molinfo $molid set frame $i

        # Get cell vectors
        set a [molinfo $molid get a]
        set b [molinfo $molid get b]
        set c [molinfo $molid get c]
        
        # We could have done it like this:
        # lassign [molinfo $molid get {a b c alpha beta gamma}] a b c alpha beta gamma

        # Check that cell info exists
        if {$a < 1.0 || $b < 1.0 || $c < 1.0} {
            puts "WARNING: Frame $i has no valid cell info (a=$a b=$b c=$c). Skipping."
            continue
        }

        # For orthorhombic box, volume = a * b * c
        # For triclinic, use full cell vectors
        set alpha [molinfo $molid get alpha]
        set beta  [molinfo $molid get beta]
        set gamma [molinfo $molid get gamma]

        # General triclinic volume formula
        set pi 3.14159265358979
        set alpha_r [expr {$alpha * $pi180}]
        set beta_r  [expr {$beta  * $pi180}]
        set gamma_r [expr {$gamma * $pi180}]

        set volume [expr {$a * $b * $c * sqrt(1.0 
            - cos($alpha_r)*cos($alpha_r) 
            - cos($beta_r)*cos($beta_r) 
            - cos($gamma_r)*cos($gamma_r) 
            + 2.0*cos($alpha_r)*cos($beta_r)*cos($gamma_r))}]

        set density [expr {$total_mass * $amu_to_g / $volume}]
        puts $fp [format "%8d  %12.6f  %14.2f  %10.4f  %10.4f  %10.4f" \
                      $i $density $volume $a $b $c]

        set sum  [expr {$sum + $density}]
    }


    close $fp
    $sel delete
    
    # Print summary
    set avg [expr {$sum / $nframes}]
    puts "--------------------------------------"
    puts "Frames analyzed: $nframes"
    puts [format "Average density: %.4f g/cm3" $avg]
    puts "Output written to: $outfile"
    puts "--------------------------------------"
}

# Usage:
# mol new system.psf
# mol addfile trajectory.dcd waitfor all
# source compute_density.tcl
# compute_density top "density.dat"