#!/usr/bin/perl -w
use strict;
#use List::MoreUtils qw / pairwise/;
use File::Temp qw/ tempfile tempdir /;
use File::Copy;
use Getopt::Long;

my $conf = {"cleanup" => 0, "oute"=> "energy", "name" => "XYZ"};
my $RUN_PREFIX = ""; #"mpirun --host node10,node11 -np 8";
my $siesta = "siesta";
my $PSEUDODIR = "/home/gkolesov/run/Pseudo/GGA";
$PSEUDODIR    = $ENV{PSEUDODIR}     if $ENV{PSEUDODIR};
$PSEUDODIR    = $ENV{SIESTA_PSEUDO} if $ENV{SIESTA_PSEUDO};

my @argv_orig = @ARGV;
GetOptions(
        "c|cleanup!"    => \$conf->{cleanup},
        "v|verbose!"    => \$conf->{verbose},
        "n|dry!"        => \$conf->{dry},
        "o|oute=s"      => \$conf->{oute},
        "x|name=s"      => \$conf->{name},
        "e|cmd=s"       => \$siesta,
        "p|ppdir=s"     => \$PSEUDODIR,
        "r|runprefix=s" => \$RUN_PREFIX,
        "h|help!"       => \&print_usage_and_die,
        );


if(@ARGV>0){
    my $remotecmd = pop @argv_orig;
    print STDERR "ARGV: @ARGV\n";
    $remotecmd  = "@ARGV" if @ARGV>1;
    #$remotecmd = "node" unless defined $remotecmd;

    exit( mysystem("$remotecmd runstruct_siesta @argv_orig") );

}


my $anumber = atom_number_table();

my ($lvec,$atoms) = parse_strout("str.out");
my $ddd =1;

# make_fdf("/tmp/y", $lvec, $atoms);
my $E=run_siesta($lvec,$atoms);
if($E==99999.){
    open(my $fw, '>', 'error') || die "Can not open file for writing: $!";
    close $fw;
}

print "E = $E\n" if $conf->{verbose};
unless($conf->{dry}){
    open(my $fw, '>', $conf->{oute}) || die "Can not open file: $!";
    print $fw $E;
    close $fw;
}

sub make_fdf{
    my ($fname,$lvec,$atoms) = @_;

    my %h=();
    my $k = 1;
    foreach my $a (@$atoms){
        $h{$a->{type}} = $k++ unless defined $h{$a->{type}};
    }

    my $acnt = scalar keys %h;

    open(my $fw, '>', $fname) || die "Can not open file for writing: $!";
    print $fw siesta_head();
    print $fw "SystemLabel $conf->{name}\n";
    print $fw "LatticeConstant 1.0 Ang\n";
    print $fw "Number_of_species $acnt\n";
    print $fw "NumberOfAtoms ",scalar @$atoms,"\n\n";
    print $fw "%block Chemical_Species_label\n";
    foreach my $t (sort {$h{$a} <=> $h{$b}} keys %h){
        printf $fw "  %-2d %-3d %s\n", $h{$t}, $anumber->{$t}, $t;
    }
    print $fw "%endblock Chemical_Species_label\n\n";
    print $fw "%block LatticeVectors\n";
    foreach my $v (@$lvec){
        printf $fw "  % -3.8f % -3.8f % -3.8f\n", @$v;
    }
    print $fw "%endblock LatticeVectors\n\n";
    
    print $fw "%block AtomicCoordinatesAndAtomicSpecies\n";
    foreach my $a (@$atoms){
        printf $fw "  % -3.8f % -3.8f % -3.8f %d\n", @{$a->{coord}},
            $h{$a->{type}};
    }
    print $fw "%endblock AtomicCoordinatesAndAtomicSpecies\n";

    close $fw;
}




sub get_energy{
    my ($siestaout) = @_;

    open(my $fh, $siestaout) || die "Can not open file: $!";
    my $readtot = 0;
    my $energy = 99999.;
    while(<$fh>){
        if($readtot){
            if(/^siesta.*?Total\s*=\s*(\S+)/i){
                $energy = $1;
                last;
            }
            next;
        }
        $readtot=1 if /^siesta:\s*Final\s+energy/;

    }
    close $fh;
    return $energy;
}

sub run_siesta{
    my ($lvec,$atoms) = @_;
    my $tdir = "siesta.$conf->{name}"; #tempdir(DIR => "./", CLEANUP => $conf->{cleanup});
    mkdir $tdir || die "Can not mkdir $tdir : $!";
    chdir $tdir || die "Can not chdir into dir : $!";

    my $fname   = "$conf->{name}.fdf";
    my $siestaout = "siesta.out";
    make_fdf($fname, $lvec, $atoms);

    my %atype = map {$_->{type} => 1}  @$atoms;
    foreach my $a (keys %atype){
        if(-e "$PSEUDODIR/$a.psf"){
            copy("$PSEUDODIR/$a.psf", "$a.psf");
        }
        else{
            die "No pseudopotential for element $a is found.\n";
        }
    }

    my $cmd = "$RUN_PREFIX $siesta <  $fname > $siestaout";

    mysystem($cmd);

    my $E = get_energy($siestaout);

    
    chdir "../" || die "Can not chdir ../ : $!";
    if($conf->{cleanup}){
        system("/bin/rm -rf $tdir");
    }

    return $E;

}


sub mysystem{
    if($conf->{dry} || $conf->{verbose}){
        print STDOUT "@_\n";
        return system("/bin/true",3) if $conf->{dry} && $conf->{verbose};
        return 0 if $conf->{dry};
    }
    
    return system(@_);
}


sub parse_strout{
    my ($file) = @_;

    open(my $fh, $file) || die "Can not open file: $!";
    my @cvec = map {$_=<$fh>; [split] } (0..2);
    my @lvec = map {$_=<$fh>; [split] } (0..2);

    my @clvec = map {cvec2cart($_,\@cvec)} @lvec;

    my @atoms = ();
    while(<$fh>){
        my @t = split /[\s,]+/;
        my $a = pop @t;
        my $v = cvec2cart(\@t,\@cvec);
        push @atoms, {type => $a, coord => $v};

    }
    close $fh;

    return (\@clvec,\@atoms);
}

sub cvec2cart{
    my ($v,$cvec) = @_;
    local $^W = 0;

    #my @t=pairwise { [map {$a *$_} @$b] } @$v,@$cvec;
    my @t = ();
    for(my $i=0; $i<@$v; $i++){
        my $a = $v->[$i];
        my $b = $cvec->[$i];
        push @t, [map {$a * $_} @$b];
    }

    my @l = ();
    for(my $i=0; $i<3; $i++){
        do { $l[$i]+=$_->[$i] } foreach @t;
    }

    return \@l;

}

sub atom_number_table{
    my %t = (
            "H"   => 1  ,
            "He"  => 2  ,
            "Li"  => 3  ,
            "Be"  => 4  ,
            "B"   => 5  ,
            "C"   => 6  ,
            "N"   => 7  ,
            "O"   => 8  ,
            "F"   => 9  ,
            "Ne"  => 10 ,
            "Na"  => 11 ,
            "Mg"  => 12 ,
            "Al"  => 13 ,
            "Si"  => 14 ,
            "P"   => 15 ,
            "S"   => 16 ,
            "Cl"  => 17 ,
            "Ar"  => 18 ,
            "K"   => 19 ,
            "Ca"  => 20 ,
            "Sc"  => 21 ,
            "Ti"  => 22 ,
            "V"   => 23 ,
            "Cr"  => 24 ,
            "Mn"  => 25 ,
            "Fe"  => 26 ,
            "Co"  => 27 ,
            "Ni"  => 28 ,
            "Cu"  => 29 ,
            "Zn"  => 30 ,
            "Ga"  => 31 ,
            "Ge"  => 32 ,
            "As"  => 33 ,
            "Se"  => 34 ,
            "Br"  => 35 ,
            "Kr"  => 36 ,
            "Rb"  => 37 ,
            "Sr"  => 38 ,
            "Y"   => 39 ,
            "Zr"  => 40 ,
            "Nb"  => 41 ,
            "Mo"  => 42 ,
            "Tc"  => 43 ,
            "Ru"  => 44 ,
            "Rh"  => 45 ,
            "Pd"  => 46 ,
            "Ag"  => 47 ,
            "Cd"  => 48 ,
            "In"  => 49 ,
            "Sn"  => 50 ,
            "Sb"  => 51 ,
            "Te"  => 52 ,
            "I"   => 53 ,
            "Xe"  => 54 ,
            "Cs"  => 55 ,
            "Ba"  => 56 ,
            "La"  => 57 ,
            "Ce"  => 58 ,
            "Pr"  => 59 ,
            "Nd"  => 60 ,
            "Pm"  => 61 ,
            "Sm"  => 62 ,
            "Eu"  => 63 ,
            "Gd"  => 64 ,
            "Tb"  => 65 ,
            "Dy"  => 66 ,
            "Ho"  => 67 ,
            "Er"  => 68 ,
            "Tm"  => 69 ,
            "Yb"  => 70 ,
            "Lu"  => 71 ,
            "Hf"  => 72 ,
            "Ta"  => 73 ,
            "W"   => 74 ,
            "Re"  => 75 ,
            "Os"  => 76 ,
            "Ir"  => 77 ,
            "Pt"  => 78 ,
            "Au"  => 79 ,
            "Hg"  => 80 ,
            "Tl"  => 81 ,
            "Pb"  => 82 ,
            "Bi"  => 83 ,
            "Po"  => 84 ,
            "At"  => 85 ,
            "Rn"  => 86 ,
            "Fr"  => 87 ,
            "Ra"  => 88 ,
            "Ac"  => 89 ,
            "Th"  => 90 ,
            "Pa"  => 91 ,
            "U"   => 92 ,
            "Np"  => 93 ,
            "Pu"  => 94 ,
            "Am"  => 95 ,
            "Cm"  => 96 ,
            "Bk"  => 97 ,
            "Cf"  => 98 ,
            "Es"  => 99 ,
            "Fm"  => 100,
            "Md"  => 101,
            "No"  => 102,
            "Lr"  => 103,
            "Rf"  => 104,
            "Db"  => 105,
            "Sg"  => 106,
            "Bh"  => 107,
            "Hs"  => 108,
            "Mt"  => 109,
            "Ds"  => 110,
            "Rg"  => 111,
            "Cn"  => 112,
            "Uut" => 113,
            "Fl"  => 114,
            "Uup" => 115,
            "Lv"  => 116,
            "Uus" => 117,
            "Uuo" => 118
    );
    return \%t;
}

sub siesta_head{
    foreach my $head ("siesta.wrap", "../siesta.wrap","../../siesta.wrap"){
	if(-f $head){
	    open(my $fh, $head) || die "Can not open $head: $!";
	    local $/ = undef;
	    my $t =  <$fh>;
	    close $fh;
	    return $t;
	}
    }

    return <<SHEAD;
DM.MixingWeight 0.01
DM.Tolerance 0.001
DM.NumberPulay 5
# DM.Require.Energy.Convergence .true.
# DM.Energy.Tolerance 0.01
PAO.BasisSize dz
SolutionMethod diagon
AtomicCoordinatesFormat Ang
MaxSCFIterations 1000
SpinPolarized False
MeshCutoff 4081.709348 eV
LatticeConstant 1.0 Ang


%block kgrid_Monkhorst_Pack
5 0 0 0.5
0 5 0 0.5
0 0 5 0.5
%endblock kgrid_Monkhorst_Pack

MD.TypeOfRun            CG          # Type of dynamics:
MD.VariableCell         .false.
#MD.NumCGsteps           50
#MD.NumCGsteps           1

WriteKbands    .false.
WriteBands     .false.
UseSaveData True
DM.UseSaveDM True
SHEAD
}

sub print_usage_and_die{
    print STDERR <<USAGE;
    $0 [OPTION] [remotecmd host]
    -v,  --verbose
    -n,  --dry          Dry mode. Doesn't actually do anything, just
                        says what it would do 
    -c,  --cleanup      Remove siesta working directory (default=false)
    -o,  --oute         File to write energy value into (default=energy)
    -x,  --name         Siesta system name label and file prefix (default=XYZ)
    -e,  --cmd          Siesta executable name (default=siesta)
    -p,  --ppdir        Pseudopotential directory; can be set also with
                        environment variables PSEUDODIR or SIESTA_PSEUDO
    -r,  --runprefix    siesta run prefix, e.g. "mpirun -np 4"
    -h,  --help         print this message

AUTHOR:
    Grigori Kolesov gkolesov@uwyo.edu

EXAMPLE:
    pollmach runstruct_siesta -e siesta.static -p ~/run/Pseudo/GGA 

    If siesta.wrap file exists in the current directory or in ../ , or in 
    ../../ the script will use it as a siesta input fdf file (and will add the
    geometry and atomic specie setup), otherwise the default at the bottom of
    this script will be used.

    Note that the siesta.wrap file should not contain LatticeConstant and
    SystemLabel fields which are also setup by the script.

    If the ~/.machines.rc file is set up and contains something like
    ssh node1  sar -u  1 1 | getvalue "Average:        all" + node -s node1
    the script should be able to run siesta via "node" script / ssh.

USAGE
    exit(0);
}