#!/usr/bin/perl -w

#  Sum stresses anually
use strict;
use lib '/pm/grflib/perl';
use GetOpt;
use FORTRAN;
my $Ni=165;
my $Nj=326;
my $Nij = $Ni*$Nj;
my @area  = ('irr','agw','asw','aco');
my @ann   = ('ppt','mi');
my @type  = ('pmp','rcg','rcs','rcc');
my @cbc   = ('sto','flo','chd','drn','et','str');
my @Type  = (@ann,@type,@cbc,@area);
my %txt   = (pmp=>'Pumping',rcg=>'GW Returns',rcs=>'SW Returns',rcc=>'Canal Leak',ppt=>'Precip',mi=>'M&I',irr=>'Irr Acres',agw=>'GW Acres',asw=>'SW Acres',aco=>'CO Acres',sto=>'Storage',chd=>'C.Head',drn=>'Springs',et=>'ET',str=>'Streams',flo=>'UnderFlow');
my @ndm   = (undef,31,28,31,30,31,30,31,31,30,31,30,31);
my %cbc   = ('STORAGE'=>'sto','CONSTANTHEAD'=>'chd','DRAINS'=>'drn','ET'=>'et','STREAMLEAKAGE'=>'str','FLOWRIGHTFACE'=>'flo','FLOWFRONTFACE'=>'flo');

my %args = GetOpt('f:o:' , f=>'../data0/cty.flg');
(@ARGV==1) || die "Usage: $0 <year>\n";
my $year  = shift @ARGV;
my @state = ('co','ks','ne');
my $flg = $args{f};
my $out = ($args{o}) ? $args{o} : "$year.cty";

#  Read flags
my (@flg) = (0);
open(FLG , "<$flg") || die "Cannot open file $flg\n";
my $line = <FLG>;
chomp $line;
my @cty = split(' ',$line);
my $wid = (@cty>9) ? 2 : 1;
for(my $i=0;$i<$Ni;$i++)
{
   my $line = <FLG>;
   chomp $line;
   for (my $j=0;$j<$Nj;$j++)
   {
      my $l = substr($line,$wid*$j,$wid);
      $flg[$i*$Nj+$j] = $l;
      ($l<0 || $l>@cty) && die "Invalid county id $l\n";
   }
}
close(FLG);

#
#  Read stress files
#
sub getdata
{
   my ($Q,$file,$opt,$mul) = @_;
   ($opt || -e $file) || die "Missing file $file\n";
   (-e $file) || return;
   open(DAT , "<$file") || die "Cannot open file $file\n";
   my @q;
   while (@q<$Nij)
   {
      my $line = <DAT>;
      defined($line) || die "Premature EOF $file ".scalar(@q)." < $Nij\n";
      chomp $line;
      push @q , split(' ',$line);
   }
   close(DAT);
   for (my $k=0;$k<$Nij;$k++)
   {
      $Q->[$flg[$k]] += $mul*$q[$k];
   }
}

#
#  Read CBC file
#
sub getcbc
{
   my ($Q,$cbc,$mo) = @_;
   my $buffer = $cbc->read() || die "Error reading CBC file for $mo\n";
   my ($step,$period,$var,$Nj,$Ni,$Nk) = unpack("L2A16L3",$buffer);
   $var =~ s/ //g;
   ($period == $mo) || die "Stress period=$period mismatch reading CBC file for $mo $var\n";
   $buffer = $cbc->read() || die "Error reading CBC file for $mo $var\n";
   my @q = unpack('d*',$buffer);

   #  Convert cubic feet per day to acre-feet for the month in two time steps
   my $typ = $cbc{$var} || die "Unknown CBC item $var\n";
   my $mul = 24*60*60*(($year%4==0 && $mo==2) ? 29 : $ndm[$mo]) / 2 / 43560;
   if ($typ eq 'flo')
   {
      my $di = ($var =~ /FRONT/) ? 1 : 0;
      my $dj = ($var =~ /RIGHT/) ? 1 : 0;
      for(my $i=0;$i<$Ni-$di;$i++)
      {
         for (my $j=0;$j<$Nj-$dj;$j++)
         {
            my $k0 = $i*$Nj+$j;
            my $k1 = ($i+$di)*$Nj+($j+$dj);
            if ($flg[$k0] != $flg[$k1])
            {
               $Q->{$typ}[$flg[$k0]] -= $mul*$q[$k0];
               $Q->{$typ}[$flg[$k1]] += $mul*$q[$k0];
            }
         }
      }
   }
   else
   {
      for (my $k=0;$k<$Nij;$k++)
      {
         $Q->{$typ}[$flg[$k]] += $mul*$q[$k];
      }
   }
}

#  Summarize data by county
my %Q;
#  Zero counters
foreach my $type (@Type)
{
   @{$Q{$type}} = (0) x (@cty+1);
}
#  Irrigated acres and precip recharge for the year
getdata($Q{ppt},"ppt/$year.rcp",0,640/12);
getdata($Q{irr},"ppt/$year.irr",1,1.0);
#  Annual M&I and acreages
foreach my $state (@state)
{
   getdata($Q{mi},"$state/$year.mi",1,-1.0);
   getdata($Q{agw},"$state/$year.agw",1,1.0);
   getdata($Q{asw},"$state/$year.asw",1,1.0);
   getdata($Q{aco},"$state/$year.aco",1,1.0);
}
#  Monthly pumping and return flows
foreach my $mo (1..12)
{
   my $date = sprintf "%7.2f" , $year+$mo/100;
   foreach my $state (@state)
   {
      foreach my $type (@type)
      {
         getdata($Q{$type},"$state/$date.$type",1,($type eq 'pmp')?-1.0:1.0);
      }
   }
}
#  Get model predicted values from CBC file
my $cbc = FORTRAN->open("<$year.cbc") || die "Cannot open CBC file $year.cbc\n";
foreach my $mo (1..12)
{
   #  Seven components, two time steps
   foreach my $nt (1 .. 14)
   {
      getcbc(\%Q,$cbc,$mo);
   }
}
#  For counties with constant heads, there is a discrepancy because
#  recharge and pumping in constant head cells are counted here, but
#  are zero in MODFLOW.  Therfore add the residual to constant head 
#  flow to balance the budget for counties with constant heads.
for (my $k=0;$k<@cty;$k++)
{
   ($Q{chd}[$k] != 0) || next;
   my $res=0;
   foreach my $type (@ann,@type,@cbc)
   {
      $res += $Q{$type}[$k];
   }
   $Q{chd}[$k] -= $res;
}

#  Output data
@cty = ('Outside',@cty);
open(OUT , ">$out") || die "Cannot open fie $out\n";
printf OUT '%-20s' , 'County';
foreach my $type (@Type)
{
   printf OUT " %12s" , $txt{$type};
}
print OUT "\n";
for (my $k=0;$k<@cty;$k++)
{
   printf OUT "%-20s" , $cty[$k];
   foreach my $type (@Type)
   {
      printf OUT " %12.1f" , $Q{$type}[$k];
   }
   print OUT "\n";
}
close(OUT);
