package skyview.survey;

import java.io.File;
import java.io.FileInputStream;
import java.io.IOException;
import nom.tam.fits.Fits;
import nom.tam.fits.Header;
import nom.tam.fits.BasicHDU;
import nom.tam.fits.BinaryTable;
import nom.tam.fits.BinaryTableHDU;
import nom.tam.fits.FitsException;
import nom.tam.util.BufferedDataInputStream;

import skyview.geometry.TransformationException;
import skyview.geometry.WCS;
import skyview.geometry.Projection;
import skyview.geometry.CoordinateSystem;
import skyview.geometry.Scaler;

import skyview.geometry.projecter.Hpx;

import skyview.executive.Settings;

import java.net.URL;
import java.nio.ByteBuffer;
import static java.nio.ByteOrder.BIG_ENDIAN;


/**
 * This class defines an image gotten by reading a HEALPix image where the
 * pixels are in the nested pixel order. This assumes the FITS structures found
 * in the WMAP data but could be adapted to other orders as needed.
 * 
 * There are several HEALPix formats supported:
 *    A single HEALPix intensity array in the primary HDU
 *    HEALPix data in the first extension HDU as a binary
 *      table with each row in the table corresponding to
 *      a single pixel.  This format can support multiple
 *      HEALPix Stokes parameter values per pixel, in particular
 *        I, Q and U and be read from the file and
 *        the derived quantities PI, PI/I, and PA can
 *        be returned if the I,Q,U are present.
 *    A single HEALPix intensity array stored in the 
 *      first extension with R rows, where each row in the table has an
 *      array of N values representing RxN pixels.
 *    A HEALPix file which has been split into multiple tiles.
 *      These tiles will be read as needed.
 */
public class HealPixImage extends Image {

    private String fitsFile;
    private Header fitsHeader;
    private Hpx hpp;

    // Note that we don't use the 'data' array here so
    // that we save space having a float rather than double array.
    private float[] values;
    private int nside;
    
    private String stokes;
        
    int       nTile;
    int       tileSize;
    int       pixLen;
    float[][] splitData;
    boolean[] readTile;
    String    prefix;
    boolean   split = false;
    int       stokesIndex;

    public HealPixImage(String file) throws SurveyException {

        setName(file);
        data = null;
        int dim = 0;

        this.fitsFile = file;
        stokes = Settings.get("STOKES", "I").toUpperCase();
        
        try {
            if (Settings.has("HEALPixArray")) {
                nside = getFromPrimaryHDU(file);
            } else if (Settings.has("BinnedHEALPix")) {
                nside = getBinnedHealPix(file);
            } else if (Settings.has("SplitHealPix")) {
                
               nside = getSplitHealPix(file);
            } else {
                // This is the default.
                nside = getRowHealPix(file);
            }

        } catch (Exception e) {
            throw new SurveyException("Unable to read file:" + file+": "+e, e);
        }
                
        try {
            setWCS(nside);
        } catch (Exception e) {
            throw new SurveyException("HEALPix WCS error: "+e, e);
        }
    }

    int getFromPrimaryHDU(String file) throws Exception {
        BasicHDU hdu = new Fits(file).readHDU();            
        values = (float[]) (hdu.getKernel());
        return hdu.getHeader().getIntValue("NSIDE");
    }
    
    BufferedDataInputStream getLocalOrRemote(String file) throws Exception {
        String[] prefixes = Settings.getArray("LocalURL");
        String local = file;
        if (file.startsWith("http")) {
            local = Util.replacePrefix(file, prefixes);
        }
        if (new File(local).exists()) {
            return new BufferedDataInputStream(new FileInputStream(local));
        } else {
            return new BufferedDataInputStream(new URL(file).openStream());
        }        
    }
    
    int getSplitHealPix(String file) throws Exception {
        
        String[] prefixes = Settings.getArray("LocalURL");
        BufferedDataInputStream bdi = getLocalOrRemote(file+".hdr1");
        Header h = Header.readHeader(bdi);
        initSplitData(file, h);
        split = true;
        return h.getIntValue("NSIDE");
    }
    
    void initSplitData(String file, Header h) {
        int pos  = file.lastIndexOf("_");
        int pow  = Integer.parseInt(file.substring(pos+1));
        nTile    = 12*(int)Math.pow(2, 2*pow);
        pixLen   = h.getIntValue("NAXIS1");
        tileSize = h.getIntValue("NAXIS2")/nTile;
        prefix   = file;
        splitData= new float[nTile][];
        readTile = new boolean[nTile];  // Defaults to false.
        if (stokes.equals("I")) {
            stokesIndex = 0;
        } else if (stokes.equals("Q")) {
            stokesIndex = 1;
        } else if (stokes.equals("U")) {
            stokesIndex = 2;
        } else {
            throw new IllegalArgumentException("Unexpected Stokes Parameter: "+stokes);
        }
    }
    
    int getBinnedHealPix(String file) throws Exception {
        BinaryTableHDU hdu = (BinaryTableHDU) new Fits(file).getHDU(1);
        float[][] data = (float[][])hdu.getColumn(0);
        int       nrow = hdu.getNRows();
        Header    hdr = hdu.getHeader();
        int       npix = hdr.getIntValue("NPIX");
        int       nside = hdr.getIntValue("NSIDE");
        if (npix == 0) {
            npix = 12*nside*nside;
        }
        int       seglen = data[0].length;
        if (nrow*seglen != npix) {
            throw new FitsException("Binned HEALPIX not valid: NPIX != NROW*SEGLEN: "+npix+" "+nrow+" "+seglen);
        }
        values = new float[npix];
        // Copy the segmented HEALPix values into a 1-d array
        int offset = 0;
        for (int n = 0; n < nrow; n += 1) {
            System.arraycopy(data[n], 0, values, offset, data[n].length);
            offset += data[n].length;
        }
        return nside;
    }
    
    int getRowHealPix(String file) throws Exception {
        BinaryTableHDU hdu = (BinaryTableHDU) new Fits(file).getHDU(1);
        String stokes = Settings.get("STOKES", "I").toUpperCase();
        // If we ask for any Stokes parameters other the I (which may
        // have been set by default, then the properly named
        // columns must be found in the table, but we don't require
        // any specific order.  If I is requested then whatever
        // is the first column is used.
        if (stokes.equals("I")) {
            // Just get intensity
            values = (float[]) hdu.getColumn(0);  // Whatever is in first column...
        } else if (stokes.equals("Q")) {
            values = (float[])hdu.getColumn("Q_STOKES");
        } else if (stokes.equals("U")) {
            values = (float[])hdu.getColumn("U_STOKES");
        }
        return hdu.getHeader().getIntValue("NSIDE");
    }
        
    void setWCS(int nside) throws Exception {
        skyview.geometry.WCS wcs;
        
        int dim = 1;
        while (Math.pow(2,dim) != nside) {
            dim += 1;
        }
        

        try {
            Projection proj = new Projection("Hpx");
            CoordinateSystem cs = CoordinateSystem.Gal;

            // The 0,0 point of the oblique projection is at (+2,+2) squares.
            // This is at the point (3 pi/4, 0) in the original projection.
            Scaler s1 = new Scaler(-3 * Math.PI / 4, 0, 1, 0, 0, 1);

            double isqrt2 = 1 / Math.sqrt(2);

            // Now rotate by 45 degrees to get into the oblique projection.
            s1 = s1.add(new Scaler(0., 0., isqrt2, isqrt2, -isqrt2, isqrt2));

            // Each square has a length of pi/sqrt(8), so pixels
            // have a length of pi/(nside*sqrt(8))
            double pixlen = Math.PI / (nside * Math.sqrt(8));

            // The oblique projection is 4x6 squares, so the
            // center is 2x3 squares from the bottom left corner.
            double crpix1 = 2 * nside;
            double crpix2 = 3 * nside;

            s1 = s1.add(new Scaler(crpix1, crpix2, 1 / pixlen, 0, 0, 1 / pixlen));

            wcs = new WCS(cs, proj, s1);

            hpp = (Hpx) proj.getProjecter();

            hpp.setOrder(dim);
            initialize(null, wcs, 4 * nside, 6 * nside, 1);

        } catch (TransformationException e) {
            throw new SurveyException("Error generating tranformation for HEALPix image", e);
        }
    }
    
    /** Get the width of the image in pixels */
    public int getWidth() {
        return 4 * nside;
    }

    /** Get the height of the image in pixels */
    public int getHeight() {
        return 6 * nside;
    }

    /** Only a scalar image.
     *  Note that if we are returning a PA image then
     *  this does not transform simply when we change coordinates.
     *  This is really a vector.
     */
    public int getDepth() {
        return 1;
    }
    

    /** Get a single element of data */
    public double getData(long ipix) {

        // convert to healpix index.
        long npix = hpp.cvtPixel(ipix);
        
        double val;
        if (split) {
            val = getSplitValue(npix);
        } else {
            if (npix < 0 || npix >= values.length) {
                val = -1;
            } else {
                val = values[(int) npix];
            }
        }

        return val;
    }

    double getSplitValue(long pix) {
        int ipix = (int) pix;
        int tile = ipix/tileSize;
        if (ipix < 0 || tile > splitData.length) {
            return -1;
        }
        if (!readTile[tile]) {
            getTile(tile);
        }
        int offset = ipix%tileSize;
        
        return splitData[tile][offset];        
    }
    
    
    void getTile(int tile) {
        String file = prefix + "."+tile;
        try {
            BufferedDataInputStream bds = getLocalOrRemote(file);
            byte[] buffer = new byte[tileSize*pixLen];
            long len = bds.read(buffer);
            if (len != tileSize*pixLen) {
                System.err.println("Error in bds read for split HEALPix");                
                throw new IOException("Bad read");
            }
            
            
            // We are assuming 4 byte floats with the first three elements
            // of the array as I,Q,U.
            splitData[tile] = parseArray(buffer, 4*stokesIndex, pixLen);
            readTile[tile] = true;
        } catch (Exception e) {
            throw new IllegalArgumentException("Unexpected error reading tile:"+file, e);
        }
        
    }
    
    float[] parseArray(byte[] buf, int offset, int delta) {
        float[] data = new float[buf.length/delta];
        ByteBuffer bb = ByteBuffer.wrap(buf).order(BIG_ENDIAN);
        for (int i=0; i<data.length; i += 1) {
            data[i] = bb.getFloat(offset);
            offset += delta;
        }
        return data;        
    }
    /**
     * Probably should happen, but just in case we support the get array
     * function.
     */
    public double[] getDataArray() {
        double[] data = new double[24 * nside * nside];

        for (int i = 0; i < data.length; i += 1) {
            data[i] = getData(i);
        }
        return data;
    }
    
    /**
     * Support changing the data! Probably won't use this...
     */
    public void setData(long npix, double val) {
        npix = (int) hpp.cvtPixel(npix);
        if (npix != -1) {
            values[(int) npix] = (float) val;
        }
    }
}
