package com.ianturton.cookbook.raster;

import java.io.File;
import java.io.IOException;
import java.util.ArrayList;
import java.util.List;

import org.geotools.coverage.GridSampleDimension;
import org.geotools.coverage.grid.GridCoverage2D;
import org.geotools.coverage.grid.io.AbstractGridCoverage2DReader;
import org.geotools.coverage.grid.io.AbstractGridFormat;
import org.geotools.coverage.grid.io.GridFormatFinder;
import org.geotools.data.Parameter;
import org.geotools.factory.CommonFactoryFinder;
import org.geotools.map.FeatureLayer;
import org.geotools.map.GridCoverageLayer;
import org.geotools.map.Layer;
import org.geotools.map.MapContent;
import org.geotools.referencing.CRS;
import org.geotools.referencing.CRS.AxisOrder;
import org.geotools.referencing.crs.DefaultGeographicCRS;
import org.geotools.styling.ChannelSelection;
import org.geotools.styling.ContrastEnhancement;
import org.geotools.styling.ContrastEnhancementImpl;
import org.geotools.styling.RasterSymbolizer;
import org.geotools.styling.SLD;
import org.geotools.styling.SelectedChannelType;
import org.geotools.styling.Style;
import org.geotools.styling.StyleFactory;
import org.geotools.swing.JMapFrame;
import org.geotools.swing.data.JParameterListWizard;
import org.geotools.swing.wizard.JWizard;
import org.geotools.util.KVP;
import org.opengis.filter.FilterFactory;
import org.opengis.referencing.FactoryException;
import org.opengis.referencing.crs.CoordinateReferenceSystem;
import org.opengis.style.ContrastMethod;

public class RasterViewer {

  private JMapFrame frame;
  private MapContent mapContent;
  private FeatureLayer layer;
  private Layer gridLayer;
  private Style style;
  private final static FilterFactory ff = CommonFactoryFinder.getFilterFactory2();
  private final static StyleFactory sf = CommonFactoryFinder.getStyleFactory();
  private static final int RED = 0;
  private static final int GREEN = 1;
  private static final int BLUE = 2;

  public RasterViewer(File file, CoordinateReferenceSystem crs) {

    frame = new JMapFrame();
    mapContent = new MapContent();

    frame.setMapContent(mapContent);
    frame.enableToolBar(true);
    frame.enableStatusBar(true);
    frame.setSize(800, 400);
    AbstractGridFormat format = GridFormatFinder.findFormat(file);
    System.out.println("got format " + format);
    AbstractGridCoverage2DReader reader = format.getReader(file);
    System.out.println("got reader " + reader);

    GridCoverage2D cov = null;
    try {
      cov = reader.read(null);
    } catch (IOException giveUp) {
      throw new RuntimeException(giveUp);
    }
    if (cov == null) {
      System.out.println("no coverage found");
      System.exit(0);
    }
    System.out.println("coverage " + cov.getProperties());

    if (crs == null) {
      crs = cov.getCoordinateReferenceSystem();
      System.err.println("Set CRS to " + crs);
    }
    if (crs == null) {
      System.err.println("No CRS defaulting to WGS84");
      crs = DefaultGeographicCRS.WGS84;
    }

    if (CRS.getAxisOrder(crs).equals(AxisOrder.EAST_NORTH)) {
      System.setProperty("org.geotools.referencing.forceXY", "true");
    }
    mapContent.getViewport().setCoordinateReferenceSystem(crs);
    Style rasterStyle = createRGBStyle(cov);

    /*
     * try { File sld = new File("raster.sld"); SLDParser stylereader = new
     * SLDParser(sf, sld.toURI().toURL()); Style[] style =
     * stylereader.readXML(); rasterStyle= style[0];
     *
     * } catch (Exception e) { e.printStackTrace(); }
     */
    Layer rasterLayer = new GridCoverageLayer(cov, rasterStyle);
    if (rasterLayer.getBounds() == null) {
      System.out.println("bad raster layer ");
      System.exit(0);
    }
    mapContent.addLayer(rasterLayer);
    frame.setVisible(true);
  }

  /**
   * This method examines the names of the sample dimensions in the provided
   * coverage looking for "red...", "green..." and "blue..." (case insensitive
   * match). If these names are not found it uses bands 1, 2, and 3 for the red,
   * green and blue channels. It then sets up a raster symbolizer and returns
   * this wrapped in a Style.
   *
   * @return a new Style object containing a raster symbolizer set up for RGB
   *         image
   */
  private Style createRGBStyle(GridCoverage2D cov) {

    // We need at least three bands to create an RGB style
    int numBands = cov.getNumSampleDimensions();
    if (numBands < 3) {
      return createGreyscaleStyle(1);
    }
    // Get the names of the bands
    String[] sampleDimensionNames = new String[numBands];
    for (int i = 0; i < numBands; i++) {
      GridSampleDimension dim = cov.getSampleDimension(i);
      sampleDimensionNames[i] = dim.getDescription().toString();
    }
    final int RED = 0, GREEN = 1, BLUE = 2;
    int[] channelNum = { -1, -1, -1 };
    // We examine the band names looking for "red...", "green...", "blue...".
    // Note that the channel numbers we record are indexed from 1, not 0.
    for (int i = 0; i < numBands; i++) {
      String name = sampleDimensionNames[i].toLowerCase();
      if (name != null) {
        if (name.matches("red.*")) {
          channelNum[RED] = i + 1;
        } else if (name.matches("green.*")) {
          channelNum[GREEN] = i + 1;
        } else if (name.matches("blue.*")) {
          channelNum[BLUE] = i + 1;
        }
      }
    }
    // If we didn't find named bands "red...", "green...", "blue..."
    // we fall back to using the first three bands in order
    if (channelNum[RED] < 0 || channelNum[GREEN] < 0 || channelNum[BLUE] < 0) {
      channelNum[RED] = 1;
      channelNum[GREEN] = 2;
      channelNum[BLUE] = 3;
    }
    // Now we create a RasterSymbolizer using the selected channels
    SelectedChannelType[] sct = new SelectedChannelType[cov.getNumSampleDimensions()];
    ContrastEnhancement ce = sf.contrastEnhancement(ff.literal(1.0), ContrastMethod.NORMALIZE);
    for (int i = 0; i < 3; i++) {
      sct[i] = sf.createSelectedChannelType(String.valueOf(channelNum[i]), ce);
    }
    RasterSymbolizer sym = sf.getDefaultRasterSymbolizer();
    ChannelSelection sel = sf.channelSelection(sct[RED], sct[GREEN], sct[BLUE]);
    sym.setChannelSelection(sel);

    return SLD.wrapSymbolizers(sym);
  }

  private Style createGreyscaleStyle(int band) {
    ContrastEnhancement ce = new ContrastEnhancementImpl();// sf.contrastEnhancement(ff.literal(1.0),
                                                           // new Normalize());
    SelectedChannelType sct = sf.createSelectedChannelType(String.valueOf(band), ce);

    RasterSymbolizer sym = sf.getDefaultRasterSymbolizer();
    ChannelSelection sel = sf.channelSelection(sct);
    sym.setChannelSelection(sel);

    return SLD.wrapSymbolizers(sym);
  }

  public static void main(String[] args) throws IOException {
    File file = null;
    String wkt = "";
    CoordinateReferenceSystem crs = null;
    if (args.length == 0) {
      // display a data store file chooser dialog for images
      List<Parameter<?>> list = new ArrayList<Parameter<?>>();
      list.add(new Parameter<File>("image", File.class, "Image", "GeoTiff or World+Image to display as basemap",
          new KVP(Parameter.EXT, "tif", Parameter.EXT, "jpg")));
      /*
       * list.add(new Parameter<File>("shape", File.class, "Shapefile",
       * "Shapefile contents to display", new KVP(Parameter.EXT, "shp")));
       */
      JParameterListWizard wizard = new JParameterListWizard("Image Lab", "Fill in the following layers", list);
      int finish = wizard.showModalDialog();

      if (finish != JWizard.FINISH) {
        System.exit(0);
      }
      file = (File) wizard.getConnectionParameters().get("image");
    } else {
      System.out.println(args[0]);
      int count = 0;
      while (args[count].startsWith("-")) {
        System.out.println("checking " + args[count]);
        if (count < args.length && args[count].endsWith("proj")) {
          count++;
          wkt = args[count];
          count++;
          continue;
        }
      }
      file = new File(args[count]);
      System.out.println(file);
      if (file == null || !file.exists()) {
        System.err.println(file + " doesn't exist");
        return;
      }

    }
    try {
      if (!wkt.isEmpty()) {

        crs = CRS.decode(wkt);

      }
    } catch (FactoryException e) {
      // TODO Auto-generated catch block
      e.printStackTrace();
    }

    System.out.println(file);
    new RasterViewer(file, crs);
  }
}