module tau_subs

  implicit none

  contains

  integer function nearest(array, value)

    ! Finds the nearest point to value in array, on the assumption array is
    ! monotoncially decreasing.

    real, dimension(:), intent(in) :: array
    real, intent(in) :: value

    integer :: num, ipos, istep, ilow, ihigh

    num=size(array,1)

    ipos=num/2
    istep = 1 + ipos/2

    do
      if (array(ipos) > value) then
        ipos=ipos-istep
      else
        ipos=ipos+istep
      end if
      ipos=max(1,ipos)
      ipos=min(num,ipos)
      istep = 1 + istep/2
      if (istep <= 2) exit
    end do
    ilow=max(1,ipos-3)
    ihigh=min(ipos+3,num)
    nearest=ilow-1+minloc(abs(array(ilow:ihigh)-value),1)

  end function nearest


  subroutine sort(xdata, ydata, ndata)

    implicit none

    real, dimension(:), intent(inout) :: xdata, ydata
    integer, intent(in) :: ndata

    integer :: i, j, l, ir
    real :: xswap, yswap

    if (ndata < 1) then
      print*, ' Error in s/r sort.  Number of data points is ', ndata
      stop
    else if (ndata > 1) then
      ! If ndata==1, then the data are already sorted!
      ir=ndata
      l=ir/2+1
      do
        if (l > 1)then
          l=l-1
          xswap=xdata(l)
          yswap=ydata(l)
        else
          xswap=xdata(ir)
          yswap=ydata(ir)
          xdata(ir)=xdata(1)
          ydata(ir)=ydata(1)
          IR=IR-1
          if (ir == 1) then
            xdata(1)=xswap
            ydata(1)=yswap
            exit
          end if
        end if
        i=l
        j=l+l
20      IF(J.LE.IR)THEN
          IF(J.LT.IR)THEN
            IF(xdata(j) < xdata(j+1)) j=j+1
          ENDIF
          IF(xswap < xdata(j))THEN
            xdata(i)=xdata(j)
            ydata(i)=ydata(j)
            I=J
            J=J+J
          ELSE
            J=IR+1
          ENDIF
          GO TO 20
        ENDIF
        xdata(i)=xswap
        ydata(i)=yswap
      end do

    end if

  end subroutine sort

end module tau_subs


program tau

  use ark_file_io
  use cmdfit_subs
  use define_star
  use tau_subs

  implicit none

  integer :: ipos, i1, i2, i, j
  real, dimension(:,:), allocatable :: data, axdata, output
  integer, dimension(2) :: naxis
  real, dimension(:), allocatable :: grad, axgrad
  integer, save :: n_grad
  integer :: iflag
  character :: flag
  character(len=10) :: sflag
  character(len=50) :: ifname

  ! For the cluster file.
  integer :: nstars, ncol, istar, jstar
  character(len=10), dimension(mcol) :: colstr
  type(a_star), dimension(:), allocatable :: star
  character(len=30) :: clus_fil

  ! Integers for colour and magnitude in pixel space.
  integer :: icol, imag, i_mag_sig, i2_low, i2_high

  ! To correct for the number of degrees of freedom.
  integer :: nparams, npts
  
  real, allocatable, dimension(:) :: prob, chi2

  ! The number of datapoints for the integration (number of simulations
  ! in the frequentist interpretation).
  integer, parameter :: n_int=10000
  real, dimension(n_int) :: prod_chi2, prod_prob

  integer :: n_val, i_val
  character(len=50) :: str_range
  real :: low, high
  integer :: jcol1, jcol2
  integer :: iminus, iband, correlated

  ! Distribution of tau for one realisation.
  integer, parameter :: n_one=10000
  real, dimension(n_one) :: tau_one, frac_one
  integer :: near

  real :: expect, sigma, harvest

  call setbug()

  print*, 'First read in the best fitting model FITS file.'
  iflag=inpark(naxis, data, axdata)
  if (iflag < 0) then
    print*, 'Error reading 2D isochrone file ', ifname
    stop
  end if

  allocate(output(naxis(1), naxis(2)))

  ! Now read the gradient file.
  call lstark_in(ifname)
  ipos=index(ifname, '.fit')
  ifname(ipos:ipos+3)='.grd'
  call nxtark_in(ifname)
  iflag=inpark(n_grad, grad, axgrad)
  if (iflag == -2) then
    print*, 'No gradient file found, using natural normalisation.'
  else if (iflag < 0) then
    print*, 'Error reading gradient file file ', ifname
    stop
  end if

  print*, '> Give the file of data points (or type end for chunk of image).'
  read(*,'(a)') clus_fil

  if (clus_fil == 'end') then

    nstars=1
    allocate(star(1))
    jcol1=1
    jcol2=2

    print*, '> Give uncertainty in magnitude.'
    read(*,*) star(1)%col(jcol1)%err

    star(1)%col(jcol2)%err=sqrt(2.0)*star(1)%col(jcol1)%err
    print*, 'Assuming uncertainty in colour of ', star(1)%col(jcol2)%err

    print*, 'And that the errors are uncorrelated.'
    correlated = 0

    print*, '> Give the number of datapoints and free parameters.'
    read(*,*) npts, nparams

  else

    call nxtcls_in(clus_fil)
    iflag=read_cluster_file(nstars, ncol, colstr, star)

    if (ncol > 2) then
      print*, 'The available magnitudes are'
      do jcol1=1, ncol
        print*, jcol1, colstr(jcol1)
      end do
      print*, 'Give the number of the magnitude you want.'
      read(*,*) jcol1
      print*, 'The available colours are'
      do jcol2=1, ncol
        print*, jcol2, colstr(jcol2)
      end do
      print*, 'Give the number of the colour you want.'
      read(*,*) jcol2
    else
      jcol1=1
      jcol2=2
    end if

    ! Check for flagged stars.
    istar=1
    do 
      if (istar > nstars) then
        exit
      else if (istar == nstars) then
        if (star(istar)%col(jcol1)%flg/='OO' &
        .or. star(istar)%col(jcol2)%flg/='OO') nstars=nstars-1
        exit
      else
        if (star(istar)%col(jcol1)%flg/='OO' &
          .or. star(istar)%col(jcol2)%flg/='OO') then
          print*, 'Removing ', star(istar)%id, star(istar)%col(jcol1)%flg, &
          star(istar)%col(jcol2)%flg
          star(istar:nstars-1)=star(istar+1:nstars)
          nstars=nstars-1
        else
          istar=istar+1
        end if
      end if
    end do

    !do istar=1, nstars
    !  print*, star(istar)%id, star(istar)%col(jcol1)%data, &
    !  star(istar)%col(jcol2)%data
    !end do
    !print*, nstars

    ! Now let's sort out whether you add or subtract the colour from the
    ! magnitude to get the other magnitude.
    iminus=index(colstr(jcol2), '-')
    if (iminus == 0) then
      print*, 'Cannot find the minus sign in colour ', colstr(jcol2)
      stop
    end if
    iband=index(colstr(jcol2), trim(colstr(jcol1)))
    if (iband == 0) then
      print*, 'Cannot find the magnitude ', trim(colstr(jcol2)), &
      ' in the colour ', trim(colstr(jcol2))
      print*, 'Assuming colours are uncorrelated.'
      correlated=0
    else if (iband > iminus) then
      print*, 'Will add ', trim(colstr(jcol1)), ' to ', trim(colstr(jcol2)), &
      ' to create ', colstr(jcol2)(1:iminus-1)
      correlated = 1
    else
      print*, 'Will subtract ', trim(colstr(jcol2)), ' from ', &
      trim(colstr(jcol1)), &
      ' to create ', trim(colstr(jcol2)(iminus+1:len(colstr(jcol2))))
      correlated = -1
    end if
    if (sum(star%col(jcol1)%err) >= 0.99*sum(star%col(jcol2)%err)) then
      print*, 'The uncertanties in magnitude are larger than in colour so'
      print*, 'I will take them as uncorrelated.'
      correlated = 0
    end if

    print*, '> Give the number of free parameters.'
    read(*,*) nparams
    npts=nstars

    ! Normalise the data.
    call natural_norm(maxval(star(1:nstars)%col(jcol1)%data,1), &
    minval(star(1:nstars)%col(jcol1)%data,1), data, axdata)    

  end if

  ! Set up the single realisation tau (x) axis.
  do i=1, n_one
    tau_one(i)=100.0*(real(i-n_one/2)-0.5)/real(n_one)
  end do
  ! And empty the fraction axis.
  frac_one=0.0

  call random_seed()

  ! Now, for each star
  each_star: do istar=1, nstars

    print*, 'Doing star ', star(istar)%id

    if (clus_fil == 'end') then

      print*, '> Give range of magnitudes to be used (<cr>=all).'
      read(*,'(a)') str_range

      if (str_range == ' ') then
        i2_low=1
        i2_high=naxis(2)
      else

        read(str_range,*) low, high
        i2_low =minloc(abs(axdata(:,2)-low ),1)
        i2_high=minloc(abs(axdata(:,2)-high),1)

        if (i2_low > i2_high) then
          iflag=i2_low
          i2_low=i2_high
          i2_high=iflag
        end if

      end if

      print*, 'Range of pixels is ', i2_low, i2_high

    else

      ! Find the nearest pixel in colour-magnitude space.
      imag=minloc(abs(axdata(1:naxis(2),2)-star(istar)%col(jcol1)%data),1)
      icol=minloc(abs(axdata(1:naxis(1),1)-star(istar)%col(jcol2)%data),1)

      ! Calculate how many pixels the magnitude uncertainty represents.
      i_mag_sig=nint(star(istar)%col(jcol1)%err/&
      abs(axdata(imag,2)-axdata(imag+1,2)))

      i2_low =max(imag-3*i_mag_sig,1)
      i2_high=min(imag+3*i_mag_sig,naxis(2))
      if (i2_low<1 .or. i2_high>naxis(2)) then
        print*, 'Warning, this star falls outside the image'
        i2_low=max(1, i2_low)
        i2_high=min(i2_high, naxis(2))
      end if
    end if

    output=0.0
    do i2=i2_low, i2_high
      do i1=1, naxis(1) 
        output(i1,i2)=likelihood(data, axdata, grad, i1, i2, &
        star(istar)%col(jcol2)%err, star(istar)%col(jcol1)%err, correlated, &
        flag)
      end do
    end do
    ! Make a file of the output data to look at it.
    !iflag=makark(naxis, output, axdata)
    !stop

    n_val=0
    do i1=1, naxis(1)
      do i2=i2_low, i2_high
        if (output(i1,i2) > tiny(output(i1,i2))) n_val=n_val+1
      end do
    end do

    ! Put the 2D array into a 1D array, missing out the zero probabilty
    ! datapoints.
    allocate(prob(n_val), chi2(n_val))
    i_val=0
    prob=0.0
    chi2=0.0
    do i1=1, naxis(1)
      do i2=i2_low, i2_high
        if (output(i1,i2) > tiny(output(i1,i2))) then
          i_val=i_val+1
          prob(i_val)=output(i1,i2)
          chi2(i_val)=-2.0*log(output(i1,i2))
        end if
      end do
    end do

    if (i_val /= n_val) then
      print*, 'Problem in tau.'
      stop
    end if

    ! Normalise the probability in the area of the image chosen to one.
    prob=prob/sum(prob)

    do j=1, n_val
      near=minloc(abs(chi2(j)-tau_one),1)
      frac_one(near)=frac_one(near)+prob(j)
    end do

    call sort(chi2, prob, n_val)

    ! Make it cumulative.
    do j=n_val, 1, -1
      prob(j)=sum(prob(1:j))
    end do

    ! Choose n_int random points.
    do i=1, n_int
      call random_number(harvest)
      j=nearest(prob,harvest)
      prod_chi2(i)=prod_chi2(i)+chi2(j)
    end do

    deallocate(prob, chi2)

    if (clus_fil == 'end') iflag=makark(naxis, output, axdata)

  end do each_star

  ! Before sorting we need a dummy x-array.
  prod_prob=1.0
  call sort(prod_chi2, prod_prob, n_int)

  ! Now really fill prd_prob.
  do i=1, n_int
    prod_prob(i)=real(n_int-i+1)/real(n_int)
  end do

  ! Lets measure the variance of the distribution, by measuring the
  ! 68.2 percent confidence interval.
  sigma = (prod_chi2(minloc(abs(0.159-prod_prob),1)) - &
           prod_chi2(minloc(abs(0.841-prod_prob),1)))/2.0
  print*, 'The sigma of the distribution is ', sigma, 'which should'
  print*, 'be approximately the square root of twice the number of free '
  print*, 'parameters ', npts-nparams, 'i.e. ', sqrt(real(2*(npts-nparams)))

  ! Note this is the exectation value were there no free parameters.
  expect=prod_chi2(n_int/2)

  ! Not sure what one is supposed to do the the x-axis to allow for
  ! free parameters for the ditribution of individual points.
  open(unit=12, file='one.tau')
  write(12,'(/,/)')
  do j=1, n_one
    write(12,*) tau_one(j), sum(frac_one(j:n_one))/real(nstars)
  end do
  close(12)

  ! Now apply the scaling to allow for free parameters.
  prod_chi2 = (prod_chi2-expect)
  prod_chi2 = prod_chi2*real(npts-nparams)/real(npts)
  prod_chi2=prod_chi2+expect-real(nparams)

  open(unit=12, file='integ.tau')
  write(12,'(/,/)')
  do i=1, n_int
    write(12,*) prod_chi2(i), prod_prob(i)
  end do
  close(12)

  print*, ' '
  print*, 'Reading best-fit tau^2 value from grid.fit'
  call nxtark_in('grid.fit')
  iflag=inpark(naxis, data, axdata)
  print*, 'Pr(', minval(data), ') is ', &
  quadint(prod_chi2, prod_prob, minval(data), sflag), &
  ' for ', nparams, ' free parameters.'  
  print*, 'Reduced tau^2 is ', 1.0 + &
  ((minval(data)-expect)/real(nstars-nparams))

end program tau
