package crushtools;

import java.io.FileOutputStream;
import java.io.PrintWriter;
import java.text.DecimalFormat;

import jnum.ExtraMath;
import jnum.Unit;
import jnum.Util;
import jnum.data.Statistics;
import jnum.data.WeightedPoint;
import jnum.math.Vector2D;


public class SteppedPol {

    SkyMap R[] = new SkyMap[4];
    SkyMap T[] = new SkyMap[4];

    SkyMap N, Q, U, I, P, F;

    String corePath;


    public static void main(String[] args) {
        SteppedPol pol = new SteppedPol();
        try {
            pol.read(args[0]);

            if(args.length > 3) pol.corePath += ".rot2";
            
            pol.calcN();
            pol.calcQ();
            pol.calcU();
            pol.U.scale(-1.0);
            
            if(args.length > 1) pol.correct(pol.Q, Double.parseDouble(args[1]));
            if(args.length > 2) pol.correct(pol.U, Double.parseDouble(args[2]));
            if(args.length > 3) pol.rotate(Double.parseDouble(args[3]) * Unit.deg);
          
            pol.calcP();
            pol.calcI();
            pol.calcF(0.02);

            pol.write();
           
        }
        catch(Exception e) { e.printStackTrace(); }
    }


    public void read(String corePath) throws Exception {
        this.corePath = corePath;
        SkyMap[] maps = new SkyMap[8];

        char[] pol = new char[] { 'R', 'T' };

        DecimalFormat df = new DecimalFormat("000");

        for(int p=0, k=0; p<pol.length; p++) for(int a=0; a<=135; a+=45, k++) {
            String fileName = corePath + "." + pol[p] + "." + df.format(a) + ".fits";
            System.err.println(">>> " + fileName);
            maps[k] = new SkyMap(fileName);
           
        }

        for(SkyMap map : maps) {
            map.convolveTo(1.0 * map.instrumentFWHM);
            level(map);
        }
        
        SkyMap common = SkyMapCoadder.createCoadded(maps, maps[0].projection);  
        common.clear();
        maps = SkyMapCoadder.getIdenticalGridMaps(maps, common);

        System.arraycopy(maps, 0, R, 0, 4);
        System.arraycopy(maps, 4, T, 0, 4);
    }

    public void calcN() {
        SkyMap S1 = getSum(0, 1);
        SkyMap S2 = getSum(1, 1);

        N = S1.getSkyMap();
        N.clear();

        SkyMapDifferencer.directAdd(S1, S2, 1.0, N);

        N.scale(0.5);
        N.fileName = corePath + ".N.fits";
    }

    public void calcQ() {
        Q = getSum(0, -1);
        Q.fileName = corePath + ".Q.fits";
    }

    public void calcU() {
        U = getSum(1, -1);
        U.fileName = corePath + ".U.fits";
    }

    public void calcI() {
        I = N.getSkyMap();
        I.clear();

        for(int i=P.sizeX; --i >=0; ) for(int j=P.sizeY; --j >= 0; ) {
            if(N.flag[i][j] != 0 || P.flag[i][j] != 0) {
                I.flag[i][j] = 1;
                continue;
            }

            I.signal[i][j] = ExtraMath.hypot(N.signal[i][j], P.signal[i][j]);
            I.weight[i][j] = N.weight[i][j];
            I.count[i][j] = N.count[i][j];
            I.flag[i][j] = 0;
        }

        I.calc();
    
        I.fileName = corePath + ".I.fits";
    }
    
    public void level(SkyMap map) {
        
        WeightedPoint[] points = new WeightedPoint[map.countPoints()];
        for(int i=points.length; --i >= 0; ) points[i] = new WeightedPoint();
        
        int k = 0;
        for(int i=map.sizeX; --i >= 0; ) for(int j=map.sizeY; --j >= 0; ) if(map.flag[i][j] == 0) {
            if(map.s2n[i][j] < 3.0) points[k++] = new WeightedPoint(map.signal[i][j], map.weight[i][j]);
        }
        
        map.offset(-Statistics.Inplace.median(points, 0, k).value());
        map.calc();
    }
    

    public void calcP() {
        P = N.getSkyMap();
        P.clear();

        for(int i=P.sizeX; --i >=0; ) for(int j=P.sizeY; --j >= 0; ) {
            if(Q.flag[i][j] != 0 || U.flag[i][j] != 0) {
                P.flag[i][j] = 1;
                continue;
            }

            double q2 = Q.signal[i][j];
            double u2 = U.signal[i][j];
            q2 *= q2; u2 *= u2;

            double sigma2Q = 1.0 / Q.weight[i][j];
            double sigma2U = 1.0 / U.weight[i][j];

            // De-bias the Rice distribution
            // The following approximation is approximately correct
            // for S/N > 4 according to Simmons & Stewart 1984
            // Better approximations are numerically nasty.
            double psigma2 = (q2*sigma2Q + u2*sigma2U) / (q2 + u2);
            double pol2 = q2 + u2 - psigma2;
            if(pol2 < 0.0) pol2 = 0.0;

            // Propagate errors properly...
            P.signal[i][j] = Math.sqrt(pol2);
            P.weight[i][j] = 1.0 / psigma2;
            P.flag[i][j] = 0;
        }

        P.calc();
        P.fileName = corePath + ".P.fits";
    }

    public void calcF(double accuracy) {

        final double minw = 1.0 / (accuracy * accuracy);

        F = I.getSkyMap();
        F.clear();

        for(int i=F.sizeX; --i >=0; ) for(int j=F.sizeY; --j >= 0; ) {
            if(I.flag[i][j] != 0 || P.flag[i][j] != 0) {
                F.flag[i][j] = 1;   
                continue;
            }

            F.signal[i][j] = P.signal[i][j] / I.signal[i][j];
            if(Double.isInfinite(F.signal[i][j])) {
                F.flag[i][j] = 1;
                continue;
            }
            
            // f = a/b --> df/db = -a/b^2 * db
            // df2 = (da / b)^2 + (a db / b2)^2 = a/b * ((da/a)^2 + (db/b)^2)
            // 1 / wf = 1/(wa * b2) + a2/(wb*b4) = a2/b2 * (1/(wa*a2) + 1/(wb*b2))
            // wf = b2/a2 / (1/(wa*a2) + 1/(wb*b2))
            double p2 = P.signal[i][j];

            if(p2 == 0.0) {
                F.flag[i][j] = 1;   
                continue;
            }


            double t2 = I.signal[i][j];
            p2 *= p2; t2 *= t2;

            F.weight[i][j] = t2 / p2 / (1.0 / (p2 * P.weight[i][j]) + 1.0 / (t2 * I.weight[i][j]));
            
            // if sigma_f > accuracy than flag the datum
            if(F.weight[i][j] < minw) F.flag[i][j] = 1;  
            else {
                F.flag[i][j] = 0;  
                F.signal[i][j] -= 1.0 / Math.sqrt(F.weight[i][j]);
                if(F.signal[i][j] < 0.0) F.signal[i][j] = 0.0;
            }
        }
       
        F.calc();
        
        F.fileName = corePath + ".F.fits";
    }
    
    
    public void writeVectors(double spacing, double maxPol) throws Exception {
        SkyMap q = Q.getSkyMap();
        SkyMap u = U.getSkyMap();
        SkyMap f = F.getSkyMap();
        
        //q.convolveTo(spacing);
        //u.convolveTo(spacing);
        //f.convolveTo(spacing);
        
        String fileName = corePath + ".reg";
        PrintWriter out = new PrintWriter(new FileOutputStream(fileName));
   
        
        double delta = spacing / q.delta;
        System.err.println("### delta " + delta);
          
        double fi0 = Math.IEEEremainder(q.fracIndexOfdX(0.0), delta);
        if(fi0 < 0.0) fi0 += delta;
        
        double fj0 = Math.IEEEremainder(q.fracIndexOfdY(0.0), delta);
        if(fj0 < 0.0) fj0 += delta;
        
        System.err.println("### start " + fi0 + "," + fj0);
          
        for(double fi = fi0; fi < q.sizeX; fi += delta) for(double fj=fj0; fj < q.sizeY; fj += delta) {
            
            final int i = (int) Math.round(fi);
            final int j = (int) Math.round(fj);
            
            if(q.flag[i][j] != 0) continue;
            if(u.flag[i][j] != 0) continue;          
            if(f.flag[i][j] != 0) continue;
   
            double angle = 0.5 * Math.atan2(u.signal[i][j], q.signal[i][j]);
            
            double l = (f.signal[i][j] / maxPol);
            if(l < 0.0) l = 0.0;
            if(l > 1.0) l = 1.0;
            
            Vector2D d = new Vector2D(l * delta * Math.sqrt(2.0), 0);
            d.scale(0.5);
            d.rotate(angle + 0.5 * Math.PI);
            
            out.println("physical;line " + Util.f2.format(fi-d.x()) + " " + Util.f2.format(fj-d.y()) + " " +
                    Util.f2.format(fi+d.x()) + " " + Util.f2.format(fj+d.y()) + " # color=white");
            
        }
        
        System.err.println();
          
        out.close();
        System.err.println(" Written " + fileName);
        
    }

    private SkyMap getSum(int firstIndex, int sign) {
        SkyMap sum = R[firstIndex].getSkyMap();
        sum.clear();

        SkyMap S1 = sum.getSkyMap();
        S1.clear();

        SkyMapDifferencer.directAdd(R[firstIndex], T[firstIndex], sign, S1);
        S1.scale(0.5);

        SkyMap S2 = sum.getSkyMap();
        S2.clear();

        SkyMapDifferencer.directAdd(R[firstIndex + 2], T[firstIndex + 2], sign, S2);
        S2.scale(0.5);

        SkyMapDifferencer.directAdd(S1, S2, sign, sum);
        sum.scale(0.5);

        sum.calc();
        
        return sum;
    }

    public void correct(SkyMap pol, double factor) {
        double factor2 = factor * factor;

        for(int i=pol.sizeX; --i >= 0; ) for(int j=pol.sizeY; --j >= 0; ) {
            if(N.flag[i][j] != 0) pol.flag[i][j] = 1; 
            else if(pol.flag[i][j] == 0) {
                pol.signal[i][j] -= factor * N.signal[i][j];  
                pol.weight[i][j] = 1.0 / (1.0 / pol.weight[i][j] + factor2 / N.weight[i][j]);
            }
        }   
    }

    public void rotate(double angle) {
        double c = Math.cos(2.0 * angle);
        double s = Math.sin(2.0 * angle);

        double c2 = c*c;
        double s2 = s*s;

        for(int i=Q.sizeX; --i >= 0; ) for(int j=Q.sizeY; --j >= 0; ) {
            if(Q.flag[i][j] != 0) U.flag[i][j] = 1; 
            else if(U.flag[i][j] != 0) Q.flag[i][j] = 1;
            else {
                double q = Q.signal[i][j];
                double u = U.signal[i][j];
                double wq = Q.weight[i][j];
                double wu = U.weight[i][j];

                Q.signal[i][j] = c * q - s * u;
                Q.weight[i][j] = 1.0 / (c2 * 1.0/wq + s2 * 1.0/wu);

                U.signal[i][j] = s * q + c * u;
                U.weight[i][j] = 1.0 / (s2 * 1.0/wq + c2 * 1.0/wu);    
            }
        }   
    }

    public void write() throws Exception {
        I.write();
        Q.write();
        U.write();
        P.write();
        F.write();
        
        //writeVectors(0.5 * Q.instrumentFWHM, 0.30);
        writeVectors(2.0 * 4.020 * Unit.arcsec, 0.06);
    }

}
