# sourcecountsim.g: definition of the sourcecountssim object
##  2003dec13   LAMA Memo 803 writeup version
#
# 
#
#  #  To simulate a field @ 90 GHz:
#
#  include 'sourcecountsim.g'
#  mycounts := sourcecountsim();
#  mycounts.setsourcecounts(smin=0.001, fieldsize=50);  
#  mycounts.setupcounts();
#  mycounts.readalpha();
#  # NOTE:  if you never reseed(), every field will be different
#  #        IF you have need to reproduce a field, reseed() with
#  #        the same number and same inputs
#  mycounts.reseed(100);
#  mylist :=  mycounts.simulatesources();
#  mycounts.scalelist(mylist, newfreq=250)
#  mycounts.testsourcecount(mylist);    
#  mycounts.done();
#
#  #  To simulate a field at a different frequency:
#
#  include 'sourcecountsim.g'
#  mycounts := sourcecountsim();
#  mycounts.readalpha();
#  mycounts.setsourcecounts(smin=0.001, fieldsize=1);  
#  mycounts.importcounts(infile='CUM.COUNTS.250GHZ', freq=250);
#  mylist :=  mycounts.simulatesources();
#  mycounts.testsourcecount(mylist);    
#  mycounts.done();
#
#
#  #  To shift counts to a different frequency:
#
#  include 'sourcecountsim.g'
#  mycounts := sourcecountsim();
# #  mycounts.readalpha('alpha.dist.test1');
#   mycounts.readalpha();
#  mycounts.setsourcecounts(smin=0.001, fieldsize=3);  
#  mycounts.setupcounts();
#  mycounts.showcounts();
#  mycounts.shiftcounts(scalefreq=180, alphakink=0.5);
#  mycounts.showcounts('CUM.COUNTS.180GHZ')
#  mylist :=  mycounts.simulatesources();
#  mycounts.done();
#
# pragma include once;

include 'quanta.g'
include 'measures.g'
include 'pgplotter.g'
include 'randomnumbers.g'

#const
sourcecountsim := subsequence () {

  private := [=];

# roughly digitized integral source counts for 90 GHz
# Taken from MMA Memo 123
#
#     Source count FLUX [Jy]:
  private.counts.x := [0.0001, 0.0002, 0.0005, 0.001, 0.002, 0.004, 0.008, 0.010, 0.020, 0.040, 0.080, 
                0.100, 0.200, 0.400, 0.800, 1.000, 2.000, 4.000, 8.000, 10.00];

#     Source count NUMBERS, normalized by S^{-3/2}/4pi
  private.counts.y := [0.100, 0.150, 0.260, 0.567, 1.051, 1.696, 2.837, 3.340, 5.037, 7.356, 10.12, 
                10.96, 13.17, 14.32, 14.32, 14.03, 12.45, 10.20, 8.193, 7.565];

  private.counts.freq := 90;  #  90 GHz

#  private.alpha.x       A spectral index distribution
#  private.alpha.y       A spectral index distribution

  private.counts.fieldsize := 10.0;
  private.counts.smin := 0.0001;
  private.counts.nbins := 500;   #  reinterpolate the source onto bins this size
  private.counts.ncx := [];       # interpolated values
  private.counts.ncy := [];       # interpolated values
  private.nsmax := 100;

  private.seed  := F;
  private.rand  := randomnumbers();

  private.didsetup := F;          # keep track of whether we need to setupcounts()

  private.mypg := F;

#############################################################
#
#                  Public functions
#
#############################################################



#
#  clean up!
#
  self.done := function() {
        wider self, private;
        if (is_tool(private.mypg) ) {
	      private.mypg.done();
	}
        private.rand.done();
        val self := F;
        val public := F;
        return T;
  }



self.getprivate := function() {
   wider private;
   return ref(private);
}

 

#  Set the source count parameters
#
#  fieldsize [deg]
#  cumdistx=[]   S	-	Jy
#  cumdisty=[]   N(S) / S^(3/2) (Jy^{1.5} sr^{-1})
#  Note on cumdist:  we fit a power law between given points
#                    and extrapolate for lower and higher S
#  smin=0.050    minimum flux, Jy, of sources to consider

   self.setsourcecounts := function(fieldsize=10.0, 
		cumdistx=F, cumdisty=F, smin=0.010) {
     wider private;
     private.counts.fieldsize := fieldsize;
     if (cumdistx != F) { private.counts.x := cumdistx; }
     if (cumdisty != F) { private.counts.y := cumdisty; }
     private.counts.smin := smin;
     return T;
  }



#
#  interpolate source counts
#
self.setupcounts := function() {
  wider private;
 
# do the power law fit to the source counts
#  Counts = coeff * S^ {alpha - 3/2 }
 nfit := len( private.counts.x ) -1;
 coeff := [1:nfit]
 alpha := [1:nfit]
 for (ii in [1:nfit]) {
    alpha[ii] := ln(private.counts.y[ii+1]/ private.counts.y[ii]) /
	         ln(private.counts.x[ii+1]/private.counts.x[ii]) 
    coeff[ii] := private.counts.y[ii]/private.counts.x[ii]^alpha[ii]
 }
 # how many sources will we have?
 imin := 1;
 if ( private.counts.smin >  private.counts.x[len(private.counts.x)]) imin := len(private.counts.x); ;
 for (ii in [1:(len(private.counts.x)-1)]) {
   if ( private.counts.smin >= private.counts.x[ii] ) {
 	imin := ii;
 	break;
   }
 }
 if ( private.counts.smin < 0.0001) {  private.counts.smin := 0.0001; }
 allskyarea := 40000;
 thisarea := private.counts.fieldsize^2;
 frac := thisarea / allskyarea;
 private.nsmax := as_integer(4 * pi * frac * coeff[imin] *  private.counts.smin^(alpha[imin]-3/2) + 0.5 )
 x := array(0.0, private.nsmax);
 y := array(0.0, private.nsmax);
 s := array(0.0, private.nsmax);
 # make regridded normalized cumulative distribution
 ncmax := private.counts.nbins;

 smax := 20.0;   # maximum S considered, in Jy
 a := (smax /  private.counts.smin) ^ ( 1/(ncmax-1) );    # multiplicative scale factor for x
 ncx :=   private.counts.smin * a ^ ([0:(ncmax-1)]);      # normalized cumulative x
 ncy :=  0 * ncx;
 for ( i in [1:ncmax] ) {
   # increment multiplier if required 
   if (imin < (len(private.counts.x)-1) && ncx[i] > private.counts.x[imin + 1]) {
 	imin := imin + 1;
   }
   ncy[i] := coeff[imin] * ncx[i]^(alpha[imin]-3/2);  # there are ncy[i] many sources brighter than ncx[i]
 }
 ncy[1] := ncy[2];  # hack to improve a bug with the first (ie, low flux) value;

private.didsetup := T;
private.counts.ncx := ncx;
private.counts.ncy := ncy;
return T;

}

#
#  print the interpolated source counts
#
self.showcounts := function(outfile=' ') {
    wider private;

    if (! private.didsetup) {
      print 'You need to run setupcounts() first!';
      return F;
    }

	if (!is_tool(private.mypg)) {
	  private.mypg := ref pgplotter();
	}
	private.mypg.page();
	x := log(private.counts.ncx);
	y := log(private.counts.ncy * private.counts.ncx^(3/2) );
	private.mypg.env( (min(x)-1), (max(x)+1),
			  (min(y)-1), (max(y)+1),
			  0, 30)
	private.mypg.line(x, y);
	private.mypg.pt( log(private.counts.x), log(private.counts.y), 10);
	private.mypg.lab( 'Source Flux [Jy]', 's^{3/2} n(S)', 'Integral Source Counts');

    print 'NOTE: the line is the private.counts.ncy, the points are private.counts.y';


   if (outfile != ' ') {
     outpipe := open(paste("> ", outfile));

     for (i in ind(private.counts.ncx)) {
       write(outpipe, private.counts.ncx[i],  private.counts.ncy[i], sep=' ');
       write(outpipe);
     }
     outpipe := F;
     print 'Look in file', outfile, 'for source counts info';
   }
   return T;
}


self.importcounts := function(infile=' ', freq=90){
  wider private;

  nlines := as_integer(split(  shell(paste(' wc -l ', infile) ) )[1]);
  inpipe :=  open(paste("< ", infile));
  private.counts.ncx := 0 * [1:nlines];
  private.counts.ncy := 0 * [1:nlines];
  private.counts.freq := freq;

  for (i in ind(private.counts.ncx)) {
       aa := split( read(inpipe) );
       private.counts.ncx[i] := as_float(aa[1]);
       private.counts.ncy[i] := as_float(aa[2]);
  }
  inpipe := F;
  print 'Read in source counts from ',infile, ' for freq= ', freq;
  private.didsetup := T;
  return T
}




#
#   read in a spectral index distribution
#
self.readalpha := function(infile='alpha.dist3') {
   wider private;
   nlines := as_integer(split(  shell(paste(' wc -l ', infile) ) )[1]);
   private.alpha := [=];
   private.alpha.x := [1:nlines] * 0;
   private.alpha.y := [1:nlines] * 0;

   inpipe := open( paste(" < ", infile));
   for (i in ind(private.alpha.x)) {
       aa := split( read(inpipe) );
       private.alpha.x[i] := as_float(aa[1]);
       private.alpha.y[i] := as_float(aa[2]);
   }
   inpipe := F;
# check normalization
   ysum := sum( private.alpha.y );
   private.alpha.y := private.alpha.y / ysum;

   private.alpha.ycum := private.cumulate( private.alpha.y );

   print 'Read in spectral index info from ',infile;
   print 'Spectral index ranges from', private.alpha.x[1], 'to',
          private.alpha.x[len(private.alpha.x)];

   return T;
}



#
#  We need to play with this an make sure it is correct!
#  Try using some delta function spectral index distributions!
#  ALSO, look for some things in setupcounts() that need to be
#  revisited by this shift function.
#
#  LAST:  write some code to WRITE/READ these other source counts
#
# 
#  shift the source counts to the next frequency
#
self.shiftcounts := function( scalefreq=300, alphakink=0.5 ) {
   wider private;

   if (! is_record( private.alpha )) {
      print 'You need to run readalpha() first!';
      return F;
   }

   if (! private.didsetup) {
      print 'You need to run setupcounts() first!';
      return F;
   }

   print 'We are steepening alpha by ', alphakink;

   freq0 := private.counts.freq;
   freqrat := scalefreq / freq0;

   print 'We will shift the source counts from', freq0, 'to', scalefreq;
   print 'Ratio = ', freqrat;

   counts0 := private.counts.ncy  * private.counts.ncx^(3/2);
   s0      := private.counts.ncx;

   buildup := private.counts.ncy * 0;
   ncounts := len(buildup);
   nalpha  := len(private.alpha.y);
   
   for (ialpha in ind(private.alpha.y)) {
     freqratalpha := freqrat^(private.alpha.x[ialpha] + alphakink);

# find boundaries
     iflux1 := ncounts;
     iflux2 := 1;
     smin := s0[1];
     smax := s0[ncounts]
     for (iflux in [1:ncounts]) {
       snew := s0[iflux] * freqratalpha;
       if (snew < smax) {
          iflux2 := iflux;
       }
     }     
     for (iflux in [ncounts:1]) {
       snew := s0[iflux] * freqratalpha;
       if (snew > smin) {
          iflux1 := iflux;
       }
     }     
     print 'Working: alpha = ', (private.alpha.x[ialpha]+alphakink), 'Using iflux ranging: ', iflux1, iflux2;

     for (iflux in [iflux1:iflux2]) {

        snew := s0[iflux] * freqratalpha;   #  this is the flux level from which we are "dipping"
        isnew := -1;
        for (inew in [1:(ncounts-1)]) {
            if (snew >= s0[inew] && snew < s0[inew+1]) {
                isnew := inew;
                break;
            }
        }
        if (isnew == -1) {
          print 'Hey, a BIG ERROR has occurred!'
          return F;
        }

        buildup[iflux] +:= counts0[ isnew ] * private.alpha.y[ialpha] /  freqratalpha;
     }
   }
   
   getminimum := buildup[ (buildup > 0) ];
   bmin := sort( getminimum )[1];
   getminimum := F;
   buildup[(buildup <= 0)] := bmin;

   private.counts.ncy := buildup / private.counts.ncx^(3/2);


   private.counts.freq := scalefreq;

   return T;
}




#
#  simulate fast switching calibration sources, including ALPHA
#
#  return record.x[], record.y[], record.s[]
#
self.simulatesources := function() {
  wider private;

  if (! private.didsetup) {
      print 'You need to run setupcounts() first!';
      return F;
  }

  sourcelist := [=];

# psum := sum(private.counts.ncy);
# private.counts.ncy := private.counts.ncy / psum;
# ncyc := private.cumulate( private.counts.ncy );    # cumulative probability
# ncyc[ncmax] := 1.0; 
#

 ncmax := private.counts.nbins;
 smax := 20;
 a := (smax /  private.counts.smin) ^ ( 1/(ncmax-1) );    # multiplicative scale factor for x


# differential:
 ncyd := 0 * private.counts.ncy;
 ncyd[2:ncmax] := private.counts.ncy[1:(ncmax-1)] - private.counts.ncy[2:ncmax];
 ncyc := private.cumulate( ncyd )
 ncyc := ncyc/ncyc[ncmax]
 for (ii in [1:private.nsmax]) {
    x[ii] := (private.random1() - 0.5) * private.counts.fieldsize;
    y[ii] := (private.random1() - 0.5) * private.counts.fieldsize;
    rand := private.random1();
    s[ii] := private.counts.ncx[len(ncyc[(ncyc < rand)])];  
# "dither" the flux level
    rand := private.random1();
    s[ii] := s[ii] * (1 + a * (rand - 0.5));
 }

 orders := order(-s);

 sourcelist.s :=  s[orders] ;
 sourcelist.x :=  x[orders] ;
 sourcelist.y :=  y[orders] ;
 sourcelist.alpha := self.pickalpha(n = len(s));
 sourcelist.freq := private.counts.freq;

 return ref sourcelist;
}


#
#  scale a list to a new frequency
#
self.scalelist := function(ref sourcelist=[=], newfreq=90) {

  snew := sourcelist.s * (newfreq / sourcelist.freq)^(-sourcelist.alpha);
  sourcelist.freq := newfreq;
  print 'NOTE: reference symmantecs changes the source list in place.';
  return T;
}

#
# test out source counts
#
  self.testsourcecount := function( sourcelist=[=] ) {
        wider private;

	countx := -sort( -sourcelist.s );
	county := [1:len(countx)] * countx^(3/2) * (1/(4*pi)) * 40000 / private.counts.fieldsize^2;

	if (private.mypg == F) {
	  private.mypg := ref pgplotter();
	}
	private.mypg.page();
	x := log(countx);
	y := log(county);
	private.mypg.env( min(x), max(x),
			  min(y), max(y),
			  0, 30)
	private.mypg.line(x, y);
	private.mypg.pt( log(private.counts.x), log(private.counts.y), 10);
	private.mypg.lab( 'Source Flux [Jy]', 's^{3/2} n(S)', 'Integral Source Counts')
  }



 self.pickalpha := function(n=1) {
   wider private;

   if (! is_record( private.alpha )) {
      print 'You need to run readalpha() first!';
      return F;
   }

   randoms := private.randomn(n);
   alphas  := 0 * randoms;

   for (i in ind(randoms)) {
     for (j in ind(private.alpha.x)) {
       if (randoms[i] < private.alpha.ycum[j]) {
           alphas[i] := private.alpha.x[j];
           break;
       }
     }
   }
   return alphas;
 }


#  HEY, to UNSEED, call  reseed(F)
 self.reseed := function(seed=100) {
   wider private;

   private.seed := seed;
   private.rand.reseed( private.seed ); 

   return T;
 }


####################################################################################
#
#       private functions
#
####################################################################################



  private.random1 := function() {
	if (is_boolean(private.seed)) {
	   return random()/2^31;
	} else {
	   return private.rand.uniform(0,1,1);
	}
  }

  private.randomn := function(n=1) {
	if (is_boolean(private.seed)) {
	   return random(n)/2^31;
	} else {
	   return private.rand.uniform(0,1,n);
	}

  }

  private.cumulate := function( ref prob=[]) {
    cumprob := prob;
    cumprob[1] := 0;
    n := len(prob);
    for (i in [2:n]) {
      cumprob[i] := cumprob[i-1] + prob[i-1];
    }
    return ref cumprob;
  }
  
}