package pvutils;
use drsutils;
use FileHandle;
use tfm;

# Dinoj Surendran dinoj@cs.uchicago.edu / dinojs@gmail.com
# Jan 2007

use constant PI    => 4 * atan2(1, 1);

sub help {
  print STDERR <<EOF;
Usage for pvutils.pm functions:

  vd2roll : finds roll given Partiview jump parameters x y z rx ry rz (needs more testing)

  speckfilename : adds .speck to input string if not present already
  makecmap : creates CMAP file for N random numbers 
  arccos, arcsin : work with radians
  getDataVarNames : given .speck file, returns list and hash of datavariables
  statsSpeck : finds number of points and columns (including x y z columns) in a speck file

  subSampleSpeck : creates speck file with random fraction of data points of input speck file
  subColsSpeck : creates speck file with only selected columns of input speck file
  subSampleColsSpeck : creates speck file with only selected columns of random fraction of input speck file

  line2speckFH : outputs a line, with optional colour and width, to a speck file specified by a filehandle
  mesh2speckFH : outputs a mesh (ie quadrogon), with optional color, to a speck file specified by a filehandle
  path2speckFH : outputs a path with N points and N-1 edges given by a Nx3 matrix, to a speck file specified by a filehandle
EOF
}


sub speckfilename{
# adds .speck to input string if not present already
# eg input:  blah       or blah.speck or blah.SPECK
# eg output: blah.speck ,  blah.speck ,  blah.SPECK
  my ($filename) = @_;
  if (($filename !~ m/speck$/) && ($filename !~ m/SPECK$/)){
    $filename = $filename.".speck";
  }
  return $filename;
}

sub makecmap{
# creates Partiview-format cmap file with random colours
# eg input: example.cmap, 100
# eg output: example.cmap is created with 101 on first line
#    and 101 lines afterwards each with a random RGB triplet
#    so example.cmap has colours 0 to 100

  my ($cmapfilename,$N) = @_;
  if (($cmapfilename !~ m/cmap$/) && ($cmapfilename !~ m/cmap$/))
    {$cmapfilename .= ".cmap";}
  open CMAP,">$cmapfilename";
  printf CMAP "%d\n",$N+1;
  for (my $i=0; $i<=$N; $i++){
    for (my $j=0; $j<3; $j++){
      $r = rand();
      printf CMAP "%0.6f ",$r;
    }
    print CMAP "\n";
  }
  close CMAP;
}



sub arccos {my $z = $_[0]; my $zsq = $z*$z; atan2(($zsq <= 1 ? sqrt(1-$z*$z) : 0), $z)}
sub arcsin {my $z = $_[0]; my $zsq = $z*$z; atan2($z, ($zsq <= 1 ? sqrt(1-$z*$z) : 0))}



sub getDataVarNames{
# gets names of variables in speck file
# input: speckfile name (includes .speck extension)
# output: ($aref,$href) where
#   @a = @{$aref} is an array with $a[$i] = name of all i-th variables, separated by " "
#        or  $a[$i] = "" if there is no line of the form "datavar $i blah" in the speck file
#        If all columns are named and there are 3+C columsn in the speck file (remember that the 
#          first three columns are x y z) then @a has C entries $a[0..$C-1].
#   %h = %{$href} is a hash that is the reverse of @a e.g. if "datavar 10 yoyo" is in speck file
#        then $h{"yoyo"} returns 10.
 
  my ($inspeck) = @_;
  open IN,$inspeck or die ("cannot open $inspeck \n");
  my @parts;
  my %datavars;
  my $nmax=-1;
  while (<IN>){
    chomp;
    @parts = split /\s+/;
    while ($#parts >= 0 && ! length($parts[0]))
      {shift(@parts);}
    if ($#parts>=2 && $parts[0] eq "datavar" && drsutils::isInt($parts[1])){
      $datavars{$parts[2]} = $parts[1];
      $nmax = ($parts[1] > $nmax ? $parts[1] : $nmax);
    }elsif (($#parts>=0) && drsutils::isNumeric($parts[0])){
      close IN;
    }
  }
  close IN;
  my @dv = ("") x ($nmax+1);

  foreach $dvname(keys %datavars){
    my $index = $datavars{$dvname};
    if (length($dv[$index])){
      $dv[$index] .= " $dvname";
    }else{
      $dv[$index] =  $dvname;
    }
  }

  return (\@dv,\%datavars);
}

sub statsSpeck{
  # input : name of speck file, including .speck extension
  # output: (N,C) 
  #    N=number of data points in speck file
  #    C=number of columns in speck file
  #
  # Assumptions: Any line whose first nonempty entry is a number is a data point
  # Once one line with a data point is found, any line afterwards that contains a number is a data point.
  # All data point lines have the same number of columns (i.e. same # nonempty entries separated by whitespace)

  my ($inspeck) = @_;
  open IN,$inspeck or die ("cannot open $inspeck \n");
  
  my $numpts = 0;
  my $numcols = 0;
  my $linecount = 0;
  while ($line = <IN>){
    chomp $line;
    if ($numcols){
      $numpts++ if ($line =~ m/\d+/);
    }else{
      my @parts = split /\s+/, $line;
      if ($#parts>=0){                              #    if (($#parts>=0) && (drsutils::isNumeric(@parts))){
	shift(@parts) if ($parts[0] eq "");
	if ($#parts>=0 && drsutils::isNumeric($parts[0])){
	  $numpts++;
	  $numcols = $#parts+1;
	}
      }
    }
  }
  close(IN);
  return ($numpts,$numcols);
}

sub subSampleSpeck{
  # input: name of speck file (including .speck), 
  #        name of output speck file (including .speck), 
  #        fraction of datapoints to keep (real number between 0 and 1)
  #        E : verbosity control (default 1000)
  #             E>0 if you want a message when every E lines are processed
  #             E=0 if you want no messages (just the start and end messages)
  #             E=-1 if you want no messages at all)
  # output: creates output speck file with a random fraction of datapoints
  #        output speck file has same headings as input speck file   
  #
  # Assumptions : same as statsSpeck function
  
  my ($inspeck,$outspeck,$frac2keep) = @_;
  my $EVERY = ($#_>3 ? $_[3] : 1000);

  ($N,$D) = statsSpeck($inspeck);
  if ($EVERY>-1){
    printf "$inspeck has $N points, each point has %d features in addition to a 3D location\n",$D-3;
    printf "$outspeck should have a fraction %s of them.\n\n",drsutils::num2string($frac2keep);
  }

  my @parts;
  my $r,$line;
  my $numKept=0;
  my $numDiscarded=0;
  my $numLinesRead=0;

  open IN,$inspeck or die ("cannot open $inspeck \n");
  open OUT, ">$outspeck";

  my $numcols = 0; # becomes positive as soon as a line with a data point is seen
  while ($line = <IN>){
    chomp $line;
    $numLinesRead++; 
    print "Processed $numLinesRead out of $N lines of $inspeck\n" if ($EVERY>0 && $numLinesRead && ! ($numLinesRead % $EVERY));

    my $ISDATAPOINTLINE = ($numcols && $line =~ m/\d+/);
    if (! $ISDATAPOINTLINE){
      my @parts = split /\s+/, $line;
      shift(@parts) if ($#parts>=0 && $parts[0] eq "");
      if ($#parts >= 0 && drsutils::isNumeric($parts[0])){
	$numcols = $#parts+1;
	$ISDATAPOINTLINE = 1;
      }	
    }

    if ($ISDATAPOINTLINE){
      if (rand() <= $frac2keep){
	print OUT "$line\n" ;
	$numKept++;
      }else{
	$numDiscarded++;
      }
    }else{
      print OUT "$line\n";
    }
  }

  if ($EVERY > -1){
    if ($frac2keep<1){
      printf "\n$outspeck has %d out of %d data points\n Fraction Requested to Keep = %s\n Fraction Actually Kept = %s\n",$numKept,$N, $frac2keep,drsutils::num2string(1.0*$numKept/$N);
    }
  }
  close IN;
  close OUT;
}

sub subColsSpeck{
  # input: name of speck file (including .speck), 
  #        name of output speck file (including .speck), 
  #        c1,...,cN : each ci is an integer between 0 and C-1 if
  #          input speck file has 3+C columns (NOTE: DO NOT COUNT THE FIRST 3 COLUMNS!)
  # output: creates output speck file with columns c1,...,cN of input speck file   
  #        datavars are modified appropriately
  #        the output speck file has the first 3 columsn (x y z) of the input speck file
  #
  # Assumptions : same as statsSpeck function

  my ($inspeck,$outspeck,@cols2keep) = @_;
  subSampleColsSpeck($inspeck,$outspeck,1,@cols2keep);
}

sub subSampleColsSpeck{
  # input: name of speck file (including .speck), 
  #        name of output speck file (including .speck), 
  #        fraction between 0 and 1,
  #        c1,...,cN : each ci is an integer between 0 and C-1 if
  #          input speck file has 3+C columns (NOTE: DO NOT COUNT THE FIRST 3 COLUMNS!)
  # output: creates output speck file with columns c1,...,cN of random fraction of input speck file   
  #        datavars are modified appropriately
  #        the output speck file has the first 3 columns (x y z) of the input speck file
  #
  # assumptions: as with statsSpeck

  my ($inspeck,$outspeck,$frac2keep,@cols2keep) = @_;   # cols2keep can be integer or string
  if ($#cols2keep < 0){
    return subSampleSpeck($inspeck,$outspeck,$frac2keep);
  }	

  my $i,$j,$EVERY = 1000;        # EVERY is verbosity control ... works like subSampleSpeck even if I cant decide yet how to get the user to input it!
  my ($N,$D) = statsSpeck($inspeck);
  my ($dvref,$dvhref) = getDataVarNames($inspeck);
  my @dv = @{$dvref};                    # currently assumes that no  datavariable has more than one name
  my %dvh = %{$dvhref};

  # make sure @cols2keep has integers from 0 to C-1 if inspeck has 3+C columns

  for $i(0..$#cols2keep){
    $j = $cols2keep[$i];
    if (drsutils::isInt($j)){
      if ($j > $#dv)
	{die (sprintf("cols2keep[%d]=%d but datavar array dv has only %d elements\n  (remember: columns of speck file are x y z 0 1 2 .. )\n",$i,$cols2keep[$i],$#dv+1));}
    }else{
      if (exists $dvh{$j}){
	$cols2keep[$i] = $dvh{$j};
      }else{
	die (sprintf("Cannot find column named '%s' in %s\n", $j, $inspeck));
      }
    }
  }

  # now @cols2keep has integers from 0 to C-1 if inspeck has 3+C columns


  open OUT, ">$outspeck";
  map {printf OUT "datavar %d %s\n",$_,$dv[$cols2keep[$_]]}  (0..$#cols2keep);

  if ($EVERY>-1){
    printf "$inspeck has $N points, each point has %d features in addition to a 3D location\n",$D-3;
    printf "$outspeck will have %sonly the columns ",($frac2keep < 1 ? sprintf("a fraction %s of them and ",drsutils::num2string($frac2keep)) : "");
    map { printf  " %d (%s)", $cols2keep[$_], $dv[$cols2keep[$_]] } (0 .. $#cols2keep);
    print "\n";
  }

  my @parts;
  my $line,$numLinesRead=0, $numKept=0;

  open IN,$inspeck or die ("cannot open $inspeck \n");
  while ($line = <IN>){
    print "Processed $numLinesRead out of $N lines of $inspeck\n" if ($EVERY>0 && $numLinesRead && ! ($numLinesRead % $EVERY));
    $numLinesRead++;

    chomp $line;
    @parts = split /\s+/, $line;
    while (($#parts>=0) && ($parts[0] eq ""))
      {shift(@parts);}

    if ($#parts>=0 && drsutils::isNumeric($parts[0])){  # line has a data point
      if ($frac2keep == 1 || rand() <= $frac2keep){
	printf OUT "%s %s %s ", drsutils::num2string($parts[0]), drsutils::num2string($parts[1]), drsutils::num2string($parts[2]);
	map {printf OUT " %s",drsutils::num2string($parts[(3+$cols2keep[$_])]) } (0 .. $#cols2keep);
	print OUT "\n";
	$numKept++;
      }
    }elsif ($parts[0] ne "datavar"){
      print OUT "$line\n";
    }
  }
  close IN;
  close OUT;

  if ($EVERY > -1){
    if ($frac2keep<1){
      printf "\n$outspeck has %d out of %d data points and 3+%d out of 3+%d columns\n Fraction Requested to Keep = %s\n Fraction Actually Kept = %s\n",$numKept,$N, $#cols2keep+1, $D-3,  $frac2keep,drsutils::num2string(1.0*$numKept/$N);
    }else{
      printf "\n$outspeck has 3+%d out of 3+%d columns\n ", $#cols2keep+1, $D-3;
    }
  }
}



sub mesh2speckFH{
  # outputs a mesh to a speck file specified by a filehandle
  # input : OFH, @xyz0, @xyz1, @xyz2, @xyz3, [$colour]
  #  OFH is the filehandle of the output speck file
  #  Each @xyzn is a 3-dim vector representing a corner of the mesh
  #  colour is a number representing the color of the mesh
  #    colour is -1 by default which means it is not ignored in the output
  
  my ($outFH, $x0,$y0,$z0,$x1,$y1,$z1,$x2,$y2,$z2,$x3,$y3,$z3,@other) = @_;
  my $colour = ($#other >= 0 ? $other[0] : -1);

     
    
  printf $outFH "mesh ";
  printf $outFH "-c $colour " if ($colour > -1);
  printf $outFH "-s solid {\n2 2\n";
  printf $outFH "$x0 $y0 $z0\n";
  printf $outFH "$x1 $y1 $z1\n";
  printf $outFH "$x2 $y2 $z2\n";
  printf $outFH "$x3 $y3 $z3\n";
  printf $outFH "}\n";  
}

sub line2speckFH{
  # Outputs a line to a speck file specified by a filehandle.
  # Lines, if they have width, are represented by two rectangles 
  #   (both solid meshes of width w) orthogonal lengthwise
  #   i.e. a prism with a cross (+) crosssection.
  #
  # input : OFH, @xyz0, @xyz1, [colour, width]
  #  OFH is the filehandle of the output speck file
  #  @xyz0 and @xyz1 are 3-dim vectors representing the two endpoints of the line 
  #  colour is a number representing the color of the mesh
  #    colour is -1 by default which means it is not ignored in the output
  #  width is a number representing the width, if any, of the line 
  #    width is -1 by default which means the line is just a single degenerate mesh

  my ($outFH, $x0,$y0,$z0,$x1,$y1,$z1,@other) = @_;
  my $colour = ($#other >= 0 ? $other[0] : -1);
  my $width =  ($#other >= 1 ? $other[1] : -1);

  if ($width == -1){
    mesh2speckFH($outFH, $x0,$y0,$z0, $x0,$y0,$z0, $x1,$y1,$z1, $x1,$y1,$z1, @other);
  }else{
    my $w = $width;

    # draw first rectangle : first corner is at (a0,b0,c0)

    my $a0 = $x0 + $w/2;
    my $b0 = $y0 + $w/2;
    my $c0 = $z0 + $w * ($x1-$x0+$y1-$y0) / (2*($z0-$z1));
    $wcur = sqrt ( ($z0-$c0)**2 + ($y0-$b0)**2 + ($x0-$a0)**2 );  #    print "$w $wcur\n";
    $a0 = $x0 + (1.0*$w/$wcur)*($a0-$x0);    $b0 = $y0 + (1.0*$w/$wcur)*($b0-$y0);    $c0 = $z0 + (1.0*$w/$wcur)*($c0-$z0);
    my $a1 = $x0 -($a0-$x0);     my $b1 = $y0 -($b0-$y0);    my $c1 = $z0 -($c0-$z0);    # second corner
    my $a2 = $a1 + ($x1-$x0);    my $b2 = $b1 + ($y1-$y0);   my $c2 = $c1 + ($z1-$z0);    #third corner
    my $a3 = $a0 + ($x1-$x0);    my $b3 = $b0 + ($y1-$y0);   my $c3 = $c0 + ($z1-$z0);    #fourth corner
#    mesh2speckFH($outFH,$a0,$b0,$c0,$a1,$b1,$c1,$a2,$b2,$c2,$a3,$b3,$c3,@other);
    mesh2speckFH($outFH,$a0,$b0,$c0,$a1,$b1,$c1,$a3,$b3,$c3,$a2,$b2,$c2,@other);
    
    # draw second rectangle : first corner is at (p0,q0,r0)
    # (p,q,r)-(x0,y0,z0) is orthogonal to both (x1,y1,z1)-(x0,y0,z0)  and (a0,b0,c0)-(x0,y0,z0) 

    my $p0 = $x0 + ( ($y1-$y0)*($c0-$z0) - ($z1-$z0)*($b0-$y0) );
    my $q0 = $y0 + ( ($z1-$z0)*($a0-$x0) - ($x1-$x0)*($c0-$z0) );
    my $r0 = $z0 + ( ($x1-$x0)*($b0-$y0) - ($y1-$y0)*($a0-$x0) );
    $wcur = sqrt ( ($z0-$r0)**2 + ($y0-$q0)**2 + ($x0-$p0)**2 );  #  print "$w $wcur\n";
    $p0 = $x0 + ($w/$wcur)*($p0-$x0);    $q0 = $y0 + ($w/$wcur)*($q0-$y0);    $r0 = $z0 + ($w/$wcur)*($r0-$z0);
    my $p1 = $x0 -($p0-$x0);     my $q1 = $y0 -($q0-$y0);     my $r1 = $z0 -($r0-$z0);    # second corner
    my $p2 = $p1 + ($x1-$x0);    my $q2 = $q1 + ($y1-$y0);    my $r2 = $r1 + ($z1-$z0);    #third corner
    my $p3 = $p0 + ($x1-$x0);    my $q3 = $q0 + ($y1-$y0);    my $r3 = $r0 + ($z1-$z0);    #fourth corner
#    mesh2speckFH($outFH,$p0,$q0,$r0,$p1,$q1,$r1,$p2,$q2,$r2,$p3,$q3,$r3,@other);
    mesh2speckFH($outFH,$p0,$q0,$r0,$p1,$q1,$r1,$p3,$q3,$r3,$p2,$q2,$r2,@other);
  }
}

sub path2speckFH{
  # outputs a mesh to a speck file specified by a filehandle
  # input : OFH, N, $aref, [$cref, $wref]
  #  OFH is the filehandle of the output speck file
  #  @a = @{$aref} is a N x 3 array, with $a[$i]=(x,y,z) for the $i-th point, $i=0:N-1
  #   They represent N-1 lines, the n-th (n=1:N-1) line joining A(n-1,:) to A(n,:).
  #  if $cref is 0, the lines do not have colour (ie they'll take the default Partiview colour)
  #   Otherwise, @color = @{$cref} is a N-1 dim vector with $color[$i] = colour of i-th line
  #  if $wref is 0, the lines do not have width 
  #   Otherwise, @width = @{$wref} is a N-1 dim vector with $color[$i] = width of i-th line
  #
  # output : 
  #  N-1 lines are output to OFH, 

  my ($OFH, $N, $aref) = @_;
  my @a = @{$aref};
  my $cref = ($#>=3 ? $_[3] : 0);
  my $wref = ($#>=4 ? $_[4] : 0);
  my @color = ($cref ? @{$cref} : ());
  my @width = ($wref ? @{$wref} : ());
  my $n;
  
  for $n(0 .. $N-1){
    line2speck($OFH, $a[$n][0], $a[$n][1], $a[$n][2], $a[1+$n][0], $a[1+$n][1], $a[1+$n][2], (scalar(@color) ? $color[$n] : -1), (scalar(@width) ? $width[$n] : -1));
  }
}

sub vd2roll{ 
# input is (x,y,z,Rx,Ry,Rz) 
# output is roll angle corresponding to input
# not entirely sure if this is correct - it should die if it doesnt work

  my @x = lookat($_[0],$_[1],$_[2], 0,0,0, 0,1,0, 0);
  my @w2c = tfm::vd2tfm(@_);  #     w2c =  tfm(z,roll) * tfm(lookat(xyz.000.up)) = tfm(z,roll) * x
                         # so  tfm(z,-roll) * w2c = tfm(lookat(xyz.000.up)) = x if w2c * inv(x) = tfm(z,roll)
  my @zroll = &tmul( @w2c, eucinv(@x));
  die if (! drsutils::nearZero($zroll[1]+$zroll[4]));
  die if (! drsutils::nearZero($zroll[0]-$zroll[5]));
  die if (! drsutils::nearZero(arccos($zroll[0])-arcsin($zroll[1])));
  arcsin($zroll[1]);
}

1;