#!/usr/bin/env perl

=head1 NAME

scoreTrackletProbabilities - Compute tracklet probabilities

=head1 SYNOPSIS

scoreTracklets [--nn=NIGHT_NUM] [--help] [TRACKLET_IDs]

  TRACKLET_IDS : individual tracklet IDs to score
  --nn NIGHT_NUM : night number to process, or all data if unspecified
  --ssm : include Solar System model in DES manifest
  --help : display man page

=head1 DESCRIPTION

Computes a complicated, cheesy and empirical likelihood that a tracklet
is of a real object based on its morphology.

=cut

use strict;
use warnings;

use FileHandle;
use Getopt::Long;
use Pod::Usage;
use List::Util qw(min max sum);
use Math::Trig;
use POSIX qw(fmod);
use Astro::SLA;

use PS::MOPS::Constants qw(:all);
use PS::MOPS::Lib qw(:all);
use PS::MOPS::DC::Instance;
use PS::MOPS::DC::Tracklet;

my $nn;
my $survey_mode;
my $debug;
my $verbose;
my $qual;           # examine PSF_QUALITY
my @in_tids;
my $test;
GetOptions(
    'nn=i' => \$nn,
    qual => \$qual,
    verbose => \$verbose,
    'survey_mode=s' => \$survey_mode,
    debug => \$debug,
    test => \$test,
) or pod2usage(2);
@in_tids = @ARGV;
pod2usage(-msg => '--nn is required') unless ($nn || @in_tids);

my $inst = PS::MOPS::DC::Instance->new(DBNAME => undef);
my $logger = $inst->getLogger;
my $dbh = $inst->dbh;
my $sql;
my $line;
my $tid;
my $det_id;


# Scoring global, from Peter Veres's work.
#my $FAST_PX_THRESH = 5.25;
my $FAST_PX_THRESH = 6.00;
my $VERY_FAST_PX_THRESH = 7.0;
my $SLOW_VEL_THRESH = 0.1;
my $MIN_S2N = 10;
my $MIN_PSF_QUALITY = .85;

# See http://svn.pan-starrs.ifa.hawaii.edu/trac/ipp/wiki/CMF_PS1_V3
# FLAGS
my $ON_SPIKE = 0x2000_0000;
my $ON_GHOST = 0x4000_0000;

# FLAGS2
my $ON_SPIKE2 = 0x0000_0008;        # flags2
my $ON_STARCORE = 0x0000_0010;      # flags2
my $ON_BURNTOOL = 0x0000_0020;      # flags2
my $DIFF_WITH_DOUBLE = 0x0000_0002; # flags2

# Prep a statement to get morph params.
my $drsth = $dbh->prepare(<<"SQL") or die $dbh->errstr;
select 
    dr.psf_chi2 / dr.psf_dof rchi2, 
    dr.moments_xx mxx, 
    dr.moments_yy myy, 
    dr.moments_xy mxy, 
    dr.f_pos f_pos, 
    dr.psf_major pmaj, 
    dr.psf_minor pmin, 
    f.de10 xaxis, 
    (f.time_stop - f.time_start) * 86400 exptime_sec, 
    f.de10 pxspc, 
    dr.flags flags, 
    dr.flags2 flags2, 
    f.pa_deg pa_deg,
    dr.diff_r_p diff_r_p,
    dr.diff_sn_p diff_sn_p,
    dr.diff_r_m diff_r_m,
    dr.diff_sn_m diff_sn_m,
    d.s2n s2n,
    dr.psf_quality psf_quality
from fields f join detections d using(field_id) join det_rawattr_v2 dr using(det_id)
where d.det_id=?
SQL


# Testing.
if ($test) {
    print analyze_rchi2(qw(0.5 0.5 0.5)), "\n";
    print analyze_rchi2(qw(0.9 0.9 0.9)), "\n";
    print analyze_rchi2(qw(1.0 1.0 1.0)), "\n";
    print analyze_rchi2(qw(1.1 1.1 1.1)), "\n";
    print analyze_rchi2(qw(1.5 1.5 1.5)), "\n";
    print analyze_rchi2(qw(2.0 2.0 2.0)), "\n";
    exit;
}

# All tracklets.
my $tref;
if (@in_tids) {
    $tref = \@in_tids;
}
else {
    my $sm = '';
    if ($survey_mode) {
        $sm = "and survey_mode like '\%$survey_mode\%'";
    }
    print STDERR "Fetching tracklets for night $nn.\n" if $verbose;
    $tref = $dbh->selectcol_arrayref("select tracklet_id from tracklets t ignore index(classification) join fields f using (field_id) where f.nn=$nn and t.classification='N' $sm");
    $logger->info(sprintf "SCORETRACKLETPROBABILITIES: %s tracklets", scalar @{$tref});
    printf STDERR "Found %d tracklets.\n", scalar @{$tref} if $verbose;
}

print "TID NDET KNOWN VTOT RCHI2 VEL1 VEL2 S2N SPIKE GHOST QUAL BURN DIPOLE PROB\n" if $debug;

foreach my $tid (@{$tref}) {
    my $trk = modct_retrieve($inst, trackletId => $tid);
    my $prob = analyze($trk);

    if (!$debug) {
        printf "T%s: %.4f\n", $trk->trackletId, $prob if $verbose;
        $trk->probability($prob);
    }
}

$drsth->finish;
exit;


sub avg {
    return (sum @_) / @_;
}


sub analyze {
    my ($trk) = @_;
    my %raw;
    my $dets = $trk->detections;

    # Get raw attr data.
    foreach my $det (@{$dets}) {
        $drsth->execute($det->detId);
        $raw{$det->detId} = $drsth->fetchrow_hashref();
    }

    # Galactic plane handling.  Life is difficult near the galactic plane, so we need to
    # have special handling there for the burn artifacts.
    my $gb = PS::MOPS::Lib::mopslib_computeGalacticLatitude($trk->extRA, $trk->extDEC);

    my $prob_rchi2 = analyze_rchi2($trk, map { $raw{$_->detId}->{rchi2} } @{$dets});
    my $prob_vel1 = analyze_vel1($trk, $dets, \%raw);
    my $prob_vel2 = 0; # analyze_vel1($trk, $dets, \%raw);
    my $prob_s2n = analyze_s2n($trk, $dets, \%raw);
    my $prob_spike = analyze_spike($trk, $dets, \%raw);
    my $prob_ghost = analyze_ghost($trk, $dets, \%raw);
    my $prob_qual = analyze_qual($trk, $dets, \%raw);

    my $prob_burn;
    my $prob_dipole;
    if (abs($gb) < 15) {     # 15 deg of galactic plane
        # Our GP strategy will be:
        # 1. Check if both detections are on burn, regardless of orientation
        # 2. If no, then accept
        # 3. If yes, then check if chi2 score is good (>0.95?).  If yes, then recompute burn with
        #   tighter control.
        $prob_burn = analyze_burn($trk, $dets, \%raw, 20);
        if ($prob_burn) {
            if ($prob_rchi2 > 0.95) {
                $prob_burn = analyze_burn($trk, $dets, \%raw, 15);
            }
        }
        $prob_dipole = analyze_dipole($trk, $dets, \%raw, 10);       # 10 px thresh for dipoles in GP
    }
    else {
        $prob_burn = analyze_burn($trk, $dets, \%raw, 15);
        $prob_dipole = analyze_dipole($trk, $dets, \%raw, 5);        # 5 px thresh for dipoles
    }

    my $final_prob;
    if ($qual) {
        $final_prob = (max $prob_rchi2, $prob_s2n) * $prob_vel1 * (1 - $prob_qual) * (1 - $prob_spike) * (1 - $prob_burn) * (1 - $prob_ghost) * (1 - $prob_dipole);
    }
    else {
        $final_prob = (max $prob_rchi2, $prob_s2n) * $prob_vel1 * (1 - $prob_spike) * (1 - $prob_burn) * (1 - $prob_ghost) * (1 - $prob_dipole);
    }

    if ($debug) {
        my $k = 'NA';
        if ($trk->knownId) {
            ($k) = $dbh->selectrow_array('select known_name from known where known_id=?', undef, $trk->knownId);
        }
        # print "TID NDET KNOWN VTOT RCHI2 VEL1 VEL2 S2N SPIKE GHOST QUAL BURN DIPOLE PROB\n" if $debug;
        print join(' ', $trk->trackletId, scalar @{$dets}, $k, (map { sprintf "%.2f", $_ } $trk->vTot, $prob_rchi2, $prob_vel1, $prob_vel2, $prob_s2n, $prob_spike, $prob_ghost, $prob_qual, $prob_burn, $prob_dipole, $final_prob)), "\n";
    }

    return $final_prob;
}


sub analyze_rchi2 {
    # Return a PSF-ness regardless of S/N.  Allow a small GCR to contribute 
    # similarly to a single good PSF value (small 1-rchi2).
    my ($trk, @rchi2) = @_;
    my @vals = map { (1 - $_) ** 2 } @rchi2;   # compute sqr dist from 1.0
    my $max = max @vals;            # allow one outlier
    my $sum = sum(@vals);
    return 1 unless $sum;

    if (@vals > 2) {
        # GCR adjustment.
        $sum += $trk->gcr_arcsec ** 2;
#        my $denom = @vals - 1 + 1;                      # -1 for outlier, +1 for GCR adj
#        return 1 / (1 + ($sum - $max) / $denom);        # remove outlier if @vals > 2
        my $denom = @vals + 1;                      # -1 for outlier, +1 for GCR adj
        return 1 / (1 + ($sum) / $denom);        # remove outlier if @vals > 2
    }
    return 1 / (1 + $sum / @vals);
}


sub analyze_vel1 {
    # Return a score based on the tracklet having
    #   "fast motion" (large computed major axis in pixels)
    #   consistent PA
    #   PA matches tracklet PA
    my ($trk, $dref, $rref) = @_;
    my $minexp_sec = min(map { $_->{exptime_sec} } values %{$rref});
    my $pxspc = (values %{$rref})[0]->{pxspc};

    # velocity for which shortest possible exposure has expected len exceeding px thresh
    my $vel_thresh = $FAST_PX_THRESH / $minexp_sec * 86400 * $pxspc / 3600;

    if ($debug) {
        print STDERR "vel_thresh=$vel_thresh deg/day\n";
    }

    if ($trk->vTot > $vel_thresh) {
        # tracklet has "fast" motion, so create a vector that is the sum
        # of all det PAs * trail len.  This resultant vector must be in the
        # direction of the tracklet PA and within .75(?) of the expected length.
        my $trk_pa_deg = fmod($trk->posAng, 180);
        $trk_pa_deg += 180 if $trk_pa_deg < 0;      # ensure between 0, 180

        my $model_sumx = 0;
        my $model_sumy = 0;
        my $sumx = 0;
        my $sumy = 0;

        my $r;
        my $ndet = 0;
        foreach my $det (@{$dref}) {
            # For this detection, compute the avg PSF width, and a PSF major axis extent
            # for a fast detection relative to this width.
            $r = $rref->{$det->detId};
            $ndet++;
            my $avg_psf_px = 2 * ($r->{pmaj} + $r->{pmin}) / 2;
            my $min_excess_px = $FAST_PX_THRESH - $avg_psf_px;
            $model_sumx += cos(deg2rad(90 + $trk_pa_deg)) * $min_excess_px;
            $model_sumy += sin(deg2rad(90 + $trk_pa_deg)) * $min_excess_px;

            # Compute major and minor axes from moments.
            my $arg2 = 4 * $r->{mxy}**2 + ($r->{mxx} - $r->{myy})**2;
            if ($arg2 < 0) {
                $arg2 = 0;
            }

            my $maj_arg = $r->{mxx} + $r->{myy} + sqrt($arg2);
            if ($maj_arg < 0) {
                $maj_arg = 0;
            }

            my $min_arg = $r->{mxx} + $r->{myy} - sqrt($arg2);
            if ($min_arg < 0) {
                $min_arg = 0;
            }

            my $maj_px = 2 * sqrt($maj_arg / 2);
            my $min_px = 2 * sqrt($min_arg / 2);
            my $xaxis = $r->{xaxis} || 1.0;
            my $theta_deg = fmod(-90 + ($xaxis < 0 ? -1 : 1) * rad2deg(0.5 * atan2(2 * $r->{mxy}, $r->{mxx} - $r->{myy})), 180);
            $theta_deg += 180 if $theta_deg < 0;

            # Choose the angle so that it's oriented as closely as possible to the tracklet PA.  This 
            # lets us handle the bidirectional uncertainty in the detection orientation (can't tell 
            # which direction the direction is pointed +/- 180 deg).
            $theta_deg -= 180 if ($trk_pa_deg - $theta_deg < -90);
            $theta_deg += 180 if ($trk_pa_deg - $theta_deg > 90);

            # If the tracklet is "very fast", require a good match to orientation.
            if ($trk->vTot > 1.5 * $vel_thresh) {
                return 0 if abs($trk_pa_deg - $theta_deg) > 20;
            }

            $maj_px -= $avg_psf_px;
            $sumx += cos(deg2rad(90 + $theta_deg)) * $maj_px;
            $sumy += sin(deg2rad(90 + $theta_deg)) * $maj_px;
        }

        # Now we have two vectors: a model based on the tracklet velocity and
        # a real one measured from the major axis fits.

        my $model_len = sqrt($model_sumx ** 2 + $model_sumy ** 2);
        my $len = sqrt($sumx ** 2 + $sumy ** 2);
        if ($len > $model_len) {
            $sumx *= $model_len / $len;
            $sumy *= $model_len / $len;
            $len = $model_len;
        }

        if ($model_len == 0) {
            return 0;
        }
        
        my $dot = ($sumx * $model_sumx + $sumy * $model_sumy) / ($model_len ** 2);
        $dot = 1.0 if $dot > 1.0;
        $dot = 0 if $dot < 0;
        return $dot;
    }
    return 1;
}


sub analyze_vel2 {
    # Return a score based on non-fast motion but large PSF extent.
    my ($trk, $dref, $rref) = @_;
    my $maxexp_sec = max(map { $_->{exptime_sec} } values %{$rref});
    my $pxspc = (values %{$rref})[0]->{pxspc};
    my $s2n_avg = avg(map { $_->s2n } @{$dref});

    return 0;

    my $vel_thresh = $FAST_PX_THRESH / $maxexp_sec * 86400 * $pxspc / 3600;

    if ($trk->vTot < $vel_thresh) {
        my $r;
        my $sum_delta = 0;
        my $num_fast = 0;
        my $trk_pa_deg = $trk->posAng;
        foreach my $det (@{$dref}) {
            # For this detection, compute the avg PSF width, and a PSF major axis extent
            # relative to this width.
            $r = $rref->{$det->detId};

            # Compute major and minor axes from moments.
            my $arg2 = 4 * $r->{mxy}**2 + ($r->{mxx} - $r->{myy})**2;
            if ($arg2 < 0) {
                $arg2 = 0;
            }

            my $maj_arg = $r->{mxx} + $r->{myy} + sqrt($arg2);
            if ($maj_arg < 0) {
                $maj_arg = 0;
            }

            my $min_arg = $r->{mxx} + $r->{myy} - sqrt($arg2);
            if ($min_arg < 0) {
                $min_arg = 0;
            }
            next if ($maj_arg < $FAST_PX_THRESH);

            my $maj_px = 2 * sqrt($maj_arg / 2);
            my $min_px = 2 * sqrt($min_arg / 2);
            my $xaxis = $r->{xaxis} || 1.0;
            my $theta_deg = -90 + $xaxis * rad2deg(0.5 * atan2(2 * $r->{mxy}, $r->{mxx} - $r->{myy}));
            my $pa_delta = abs(($trk_pa_deg - $theta_deg) % 180);
            $sum_delta += $pa_delta;
            $num_fast++;
        }
        return 0;
    }
    return 0;
} 


sub analyze_s2n {
    # Return a score based on moderate to high S/N.  The idea here is that high, slowish S/N tracklets are
    # rare enough that we don't have to worry about morphology.  If the GCR is small then we will boost.
    my ($trk, $dref, $rref) = @_;
    my $s2n_avg = avg(map { $_->s2n } @{$dref});

    # Might want a max vel eventually, we'll see.
    my $gcr_boost;
    if ($s2n_avg <= $MIN_S2N) {
        $gcr_boost = 1.2 * $MIN_S2N / (1 + 5 * ($trk->gcr_arcsec < .05 ? 0 : ($trk->gcr_arcsec - .05))); # max boost at 0.05
    }
    else {
        $gcr_boost = 1.2 * $MIN_S2N / (1 + 5 * ($trk->gcr_arcsec < .10 ? 0 : ($trk->gcr_arcsec - .10))); # max boost at 0.05
    }
    return 1 - ($MIN_S2N / ($s2n_avg + $gcr_boost));      # check 50 => .8, 10 => 0
}


sub analyze_spike {
    # Return a score based on ON_SPIKE content.
    my ($trk, $dref, $rref) = @_;
    my $thresh = int(@{$dref} * .5 + 1.49999);      # want 2/2, 2/3, 3/4, 3/5, 4/6
    my $ct = 0;
    foreach my $r (values %{$rref}) {
        $ct++ if ($r->{flags} & $ON_SPIKE || $r->{flags2} & $ON_SPIKE2 || $r->{flags2} & $ON_STARCORE);
    }
    return $ct >= $thresh ? 1 : 0;
}


sub analyze_ghost {
    # Return a score based on ON_GHOST content.  All detections must have this set to fail.
    my ($trk, $dref, $rref) = @_;
    my $thresh = scalar @{$dref};
    my $ct = 0;
    foreach my $r (values %{$rref}) {
        $ct++ if ($r->{flags} & $ON_GHOST);
    }
    return $ct >= $thresh ? 1 : 0;
}


sub analyze_qual {
    # Return a score based on how many detections have PSF_QUALITY < some threshold.
    # High score means poor quality.
    my ($trk, $dref, $rref) = @_;
    my $thresh = scalar @{$dref};
    my $ct = 0;
    foreach my $r (values %{$rref}) {
        $ct++ if $r->{psf_quality} < $MIN_PSF_QUALITY;
    }
    return $ct >= $thresh ? 1 : 0;
}


sub analyze_dipole {
    # Return a score based on DIFF_WITH_DOUBLE and 1 < R < 5 px and matched S/N > 5 * detection S/N.
    my ($trk, $dref, $rref, $dipole_thresh_px) = @_;
    my $ct = 0;
    my $thresh;
    $dipole_thresh_px ||= 5;
    
    if (scalar @{$dref} == 2) {
        $thresh = 1;        # only need 1 of 2.
    }
    else {
        $thresh = int(@{$dref} * .5 + 1.49999);      # want 2/3, 3/4, 3/5, 4/6
    }

    foreach my $r (values %{$rref}) {   
        if (($r->{flags2} & $DIFF_WITH_DOUBLE)
            and ($r->{diff_r_p} > 0 and $r->{diff_r_p} < $dipole_thresh_px)
            and ($r->{diff_sn_p} >= 5 * $r->{s2n})
            and ($r->{diff_r_m} > 0 and $r->{diff_r_m} < $dipole_thresh_px)
            and ($r->{diff_sn_m} >= 5 * $r->{s2n})
        ) {
            $ct++;
        }
    }

    return $ct >= $thresh ? 1 : 0;
}


sub det_pa {
    # Compute major and minor axes from moments.
    my ($r) = @_;
    my $arg2 = 4 * $r->{mxy}**2 + ($r->{mxx} - $r->{myy})**2;
    if ($arg2 < 0) {
        $arg2 = 0;
    }

    my $maj_arg = $r->{mxx} + $r->{myy} + sqrt($arg2);
    if ($maj_arg < 0) {
        $maj_arg = 0;
    }

    my $min_arg = $r->{mxx} + $r->{myy} - sqrt($arg2);
    if ($min_arg < 0) {
        $min_arg = 0;
    }

    my $maj_px = 2 * sqrt($maj_arg / 2);
    my $min_px = 2 * sqrt($min_arg / 2);
    my $xaxis = $r->{xaxis} || 1.0;
    my $theta_deg = fmod(-90 + ($xaxis < 0 ? -1 : 1) * rad2deg(0.5 * atan2(2 * $r->{mxy}, $r->{mxx} - $r->{myy})), 180);
    return $theta_deg;
}


sub analyze_burn {
    # Return a score based on ON_BURTOON + PA alignment content.  If we are near the galactic plane,
    # say within 10 deg, then we need to be more aggressive at filtering tracklets where the detections
    # are on burns.
    my ($trk, $dref, $rref, $pa_thresh) = @_;
    my $ct = 0;
    my $thresh = int(@{$dref} * .5 + 1.49999);      # want 2/2, 2/3, 3/4, 3/5, 4/6
    $pa_thresh ||= 15;

    foreach my $r (values %{$rref}) {   
        my $field_pa = $r->{pa_deg};
        $field_pa += 360 if $field_pa < 0;
        $field_pa -= 360 if $field_pa >= 360;
        $field_pa -= 180 if $field_pa >= 180;

        my $det_pa = det_pa($r);
        my $pa_delta = abs(($det_pa - $field_pa) % 180);
        $pa_delta -= 180 if $pa_delta > 90;
        $pa_delta = abs($pa_delta);
        $ct++ if ($r->{flags2} & $ON_BURNTOOL and $pa_delta < $pa_thresh);
    }
    return $ct >= $thresh ? 1 : 0;
}
