#!/usr/bin/perl -w

#
#  Make Stream Package
#
#  Willem A. Schreuder
#  October 9, 2003
#

use strict;
use lib '/pm/grflib/perl';
use DBF;
use GetOpt;
use Calendar;

#
#  Segments at which flows may be set
#
my %seg =
   (
   Bonny    => 176, #  Outflow from toe drains below Bonny Reservoir
   CCP      => 122, #  Colorado Compact Compliance Pipleine
   Rock     =>  59, #  Nebraska Rock Creek Pipeline
   Medicine =>   1, #  Nebraska Medicine Creek Pipeline
   );

#
#  Redefinition of reservoirs
#
my %update =
   (
   Bonny => {#  Stage for area in square mile increments
             AREA => [3638.00,3647.51,3659.00,3670.17,3679.83],
             #  Reservoir is dry
             DRY => {
                #  Cell definitions
                CELL => ['    1  108   89  150    1    -1.00000000     1.000 1.1117000  3661.650  3671.650 Bonny',
                         '    1  107   89  150    2     0.00000000     1.000 0.9359200  3659.445  3669.445 Bonny',
                         '    1  107   90  150    3     0.00000000     1.000 3.3986000  3657.240  3667.240 Bonny',
                         '    1  107   91  150    4     0.00000000     1.000 3.2647000  3646.015  3656.015 Bonny',
                         '    1  107   92  150    5     0.00000000     1.000 0.0108170  3634.790  3644.790 Bonny',
                         '    1  106   92  150    6     0.00000000     1.000 0.0523530  3587.240  3597.240 Bonny',
                         '    1  106   92  150    7     0.00000000     1.000         0         0         0 Bonny',
                         '    1  106   92  150    8     0.00000000     1.000         0         0         0 Bonny'],
                #  Routing
                TRIB => {150=>'  140  141',156=>'  150    0'},
                },
             #  Reservoir is small (order of cells determines sequence of cells depending on stage)
             SML => {
                #  Cell definitions
                CELL => ['    1  106   91  150    1        1000000     1.000    32.267         0  3640.000 Bonny',
                         '    1  107   91  150    2     0.00000000     1.000    32.267         0  3640.000 Bonny',
                         '    1  107   90  150    3     0.00000000     1.000    32.267         0  3640.000 Bonny',
                         '    1  107   89  150    4     0.00000000     1.000    32.267         0  3640.000 Bonny',
                         '    1  107   89  150    5     0.00000000     1.000         0         0  3640.000 Bonny',
                         '    1  107   89  150    6     0.00000000     1.000         0         0  3640.000 Bonny',
                         '    1  107   89  150    7     0.00000000     1.000         0         0  3640.000 Bonny',
                         '    1  107   89  150    8     0.00000000     1.000         0         0  3640.000 Bonny'],
                #  Routing
                TRIB => {150=>'    0    0',156=>'    0    0'},
                },
             #  Reservoir is big (cells set from input file)
             BIG => {
                #  Routing
                TRIB => {150=>'    0    0',156=>'    0    0'},
                },
             }
   );

#
#  Check that tribs match
#
sub cktrib
{
   my ($res,$tp,$val) = @_;
   foreach my $type ('BIG',$tp)
   {
      exists($update{$res}{$type}{TRIB}) || die "No $type tribs specified for $res\n";
   }
   #  Check that tribs match
   my @big = sort {$a <=> $b} keys %{$update{$res}{BIG}{TRIB}};
   my @trb = sort {$a <=> $b} keys %{$update{$res}{$tp}{TRIB}};
   (@big==@trb) || die "Trib count differs $res: $tp and BIG\n";
   #  Check that tribs match
   for (my $k=0;$k<@big;$k++)
   {
      ($trb[$k]==$big[$k]) || die "Tribs mismatch $res: $tp and BIG\n";
      !$val || ($update{$res}{$tp}{TRIB}{$trb[$k]} eq $update{$res}{BIG}{TRIB}{$big[$k]}) || die "Trib values mismatch $res: $tp and BIG\n";
   }
}

#
#  Check command line arguments
#
my %args = GetOpt('f:r:q:',f=>'../data0/12s3.str',r=>'../data/reservoir.dbf',q=>'../data/flow.dbf');
my $str0 = $args{f};
(@ARGV>0) || die "Usage: $0 <year(s)>\n";

#
#  Database of reservoir levels
#
my $dbres = DBF->open($args{r},'DIRECT');
$dbres->index('MONTH','N');
#
#  Database of pipeline flows
#
my $dbq = DBF->open($args{q},'DIRECT');
$dbq->index('MONTH','N');

#
#  Snarf stream network
#
my (@STR,@WSM,@TRB,%index,%ptr);
open(STR , "<$str0") || die "Cannot open file $str0\n";
my $hdr = <STR>;
chomp $hdr;
my ($nstr,$nseg) = split(" ",$hdr);
<STR>;
for (my $l=0;$l<$nstr;$l++)
{
   my $line = <STR>;
   chomp $line;
   $STR[$l] = $line;
   #  Pointer from seg:rch to index
   my (undef,undef,undef,$seg,$rch) = split(' ',$line);
   $ptr{"$seg:$rch"} = $l;
   #  Array of reaches for this reservoir
   if (length($line)>81)
   {
      my $res = substr($line,81);
      push @{$index{$res}} , $l;
   }
}
#  Width/slope/manning
for (my $l=0;$l<@STR;$l++)
{
   my $line = <STR>;
   chomp $line;
   $WSM[$l] = $line;
}
#  Tributaries
for (my $l=0;$l<$nseg;$l++)
{
   my $line = <STR>;
   chomp $line;
   $TRB[$l] = $line;
}
close(STR);
#
#  Segments for flow locations
#
foreach my $loc (keys %seg)
{
   my $sr = "$seg{$loc}:1";
   exists($ptr{$sr}) || die "Flow $loc not found in $str0\n";
   #  Map segment to index
   $seg{$loc} = $ptr{$sr};
}
#
#  Reservoir list
#
my @res = sort keys %index;
#  Set for update
foreach my $res (@res)
{
   #  Mark reservoir
   $update{$res}{N} = @{$index{$res}};
   #  Set BIG from input data
   for (my $k=0;$k<$update{$res}{N};$k++)
   {
      $update{$res}{BIG}{CELL}[$k] = $STR[$index{$res}[$k]];
   }
   #  Set no tribs for BIG if none specified
   exists($update{$res}{BIG}{TRIB}) || ($update{$res}{BIG}{TRIB} = ());
   #  Check data for dry reservoir
   if (exists($update{$res}{DRY}))
   {
      #  Check that cell count matches
      ($update{$res}{N} == @{$update{$res}{DRY}{CELL}}) || die "Updated DRY cell count does not match $str0\n";
      #  Outflow location must be specified
      exists($seg{$res}) || die "Outflow location for $res not specified\n";
      #  Check that tribs match
      cktrib($res,'DRY',0);
   }
   #  Check data for small reservoir
   if (exists($update{$res}{SML}))
   {
      #  Check that cell count matches
      ($update{$res}{N} == @{$update{$res}{SML}{CELL}}) || die "Updated SML cell count does not match $str0\n";
      #  Make sure area is specified
      exists($update{$res}{AREA}) || die "Area not specified for $res\n";
      #  Check that tribs match
      cktrib($res,'SML',1);
   }
}

printf "$hdr\n";
foreach my $ARGV (@ARGV)
{
   my ($yr0,$yr1) = split('-',$ARGV);
   defined($yr1) || ($yr1 = $yr0);
   #
   #  Initialize reservoir stage
   #
   my %new;
   @new{@res} = $dbres->lookup('MONTH',100*($yr0-1)+12,'N',@res);

   #
   #  Generate stream file
   #
   foreach my $yr ($yr0 .. $yr1)
   {
      foreach my $mo (1 .. 12)
      {
         my $ym = 100*$yr+$mo;
         #  Acre-feet to cfs conversion factor
         my $af2cfs = 43560/(86400*DaysInMonth($yr,$mo));
         #  Get flows
         my %q;
         foreach my $flow ($dbq->select('MONTH',$ym,'N','LOC','AF'))
         {
            my ($loc,$af) = @{$flow}{'LOC','AF'};
            #  Check that the location is valid
            exists($seg{$loc}) || die "Unknown flow location $loc at $mo/$yr\n";
            #  Negative flows are prohibited
            ($af<0) && die "Negative flows are prohibited location $loc at $mo/$yr\n";
            #  Flow in cfs
            $q{$loc} = sprintf '%15.6f' , $af2cfs*$af;
         }
         #  Copy of arrays for this month
         my @str = @STR;
         my @wsm = @WSM;
         my @trb = @TRB;
         #  Update end of month elevation
         my %old = %new;
         @new{@res} = $dbres->lookup('MONTH',$ym,'N',@res);
         #  Apply reservoir data to stream network
         foreach my $res (@res)
         {
            #  This must never happen
            defined($new{$res}) || die "Missing '$res' $ym\n";
            #  Determine reservoir status
            my $type  = ($new{$res}==0) ? 'DRY' : ($new{$res}>0 ? 'BIG' : 'SML');
            #  Make sure we have data for this type
            exists($update{$res}{$type}) || die "No data for $res status $type at $mo/$yr\n";
            #  Average BOM and EOM stage if previous or current is not small or empty
            my $stage  = sprintf '%10.3f' , ($new{$res}<=0) ? abs($new{$res}) : ($old{$res}>0 ? ($old{$res}+$new{$res})/2 : $new{$res});
            #  Set routing (may be empty)
            while (my ($seg,$trb) = each %{$update{$res}{$type}{TRIB}})
            {
               $trb[$seg-1] = $trb;
            }
            #  Reservoir is empty
            if ($type eq 'DRY')
            {
               #  Set stream cells and manning
               for (my $k=0;$k<$update{$res}{N};$k++)
               {
                  my $l = $index{$res}[$k];
                  $str[$l] = $update{$res}{DRY}{CELL}[$k];
                  substr($wsm[$l],20,10) = '     0.030';
               }
               #  Set inflow to upstream outflow
               exists($q{$res}) && warn "$res outflow ignored $mo/$yr\n";
               $q{$res} = '             -1';
            }
            #  Small reservoir
            elsif ($type eq 'SML')
            {
               #  Set cell conductaces from stage/area
               my $f = 999.9;
               ($stage < $update{$res}{AREA}[0]) && die "$res stage $stage below A/C range\n";
               for (my $k=0;$k<@{$update{$res}{AREA}}-1;$k++)
               {
                  if ($update{$res}{AREA}[$k+1]>$stage)
                  {
                     $f = $k + ($stage - $update{$res}{AREA}[$k]) / ($update{$res}{AREA}[$k+1]-$update{$res}{AREA}[$k]);
                     last;
                  }
               }
               ($f>999) && die "$res stage $stage out of A/C range\n";
               #  Set stream cells
               for (my $k=0;$k<$update{$res}{N};$k++)
               {
                  my $l = $index{$res}[$k];
                  #  Set stage
                  substr($str[$l],70,10) = $stage;
                  #  Set conductance to 0 for later cells
                  if ($k > $f)
                  {
                     substr($str[$l],50,10) = '         0';
                  }
                  #  Scale conductance for partial cell
                  elsif ($k+1 > $f)
                  {
                     substr($str[$l],50,10) = sprintf '%10.3f' , ($f-$k)*substr($str[$l],50,10);
                  }
               }
            }
            #  Big reservoir
            else
            {
               #  Set stage
               for (my $k=0;$k<$update{$res}{N};$k++)
               {
                  substr($str[$index{$res}[$k]],70,10) = $stage;
               }
            }
         }
         #  Apply flows to network
         while (my ($loc,$q) = each %q)
         {
            substr($str[$seg{$loc}],25,15) = $q;
         }
         #  Write streams to file
         printf "%10d%10d%10d %2d/%4d\n" , $nstr , -1 , -1 , $mo , $yr;
         foreach my $line (@str,@wsm,@trb)
         {
            print "$line\n";
         }
      }
   }
}
