/*
 * RiskScape™ Copyright New Zealand Institute for Earth Science Limited
 * (Earth Sciences New Zealand) is distributed for research purposes only
 * under the terms of AGPLv3.
 *
 * RiskScape™ Copyright 2026 New Zealand Institute for Earth Science
 * Limited (Earth Sciences New Zealand). All rights reserved. Source code
 * available under the AGPLv3.
 * 
 * This program is free software: you can redistribute it and/or modify it under
 *  the terms of the GNU Affero General Public License as published by the Free
 *  Software Foundation, either version 3 of the License, or (at your option) any
 *  later version.
 * 
 * This program is distributed for RESEARCH PURPOSES ONLY, in the hope that it will
 * be useful for research and education initiatives.
 * 
 * If you are not a researcher, or you are a researcher who wishes to use this
 * program on terms other than AGPLv3 (including those who wish to restrict the
 * distribution of any source code created using this program), please contact:
 * https://riskscape.org.nz
 * 
 * This program is distributed WITHOUT ANY WARRANTY; without even the implied
 * warranty of MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
 * Affero General Public License for more details.  You should have received a copy
 * of the GNU Affero General Public License along with this program.  If not, see
 * <http://www.gnu.org/licenses/>.
 * 
 * By way of summary only, under the AGPLv3:
 *     • Permissions of this strongest copyleft license are conditioned
 *       on making available complete source code of licensed works and
 *       modifications, which include larger works using a licensed work,
 *       under the same license.
 *     • Copyright and license notices must be preserved.
 *     • Contributors provide an express grant of patent rights.
 *     • When a modified version is used to provide a service over a
 *       network, the complete source code of the modified version must be made
 *       available.
 */
package nz.org.riskscape.oq;

import java.util.Map;
import java.util.concurrent.atomic.AtomicLong;

import org.locationtech.jts.geom.Coordinate;
import org.locationtech.jts.geom.GeometryFactory;
import org.geotools.api.referencing.crs.CoordinateReferenceSystem;

import com.google.common.collect.ImmutableMap;
import com.google.common.collect.Range;

import lombok.Getter;
import nz.org.riskscape.engine.Engine;
import nz.org.riskscape.engine.Tuple;
import nz.org.riskscape.engine.bind.ParameterField;
import nz.org.riskscape.engine.coverage.TypedCoverage;
import nz.org.riskscape.engine.data.BaseBookmarkResolver;
import nz.org.riskscape.engine.data.BookmarkParameters;
import nz.org.riskscape.engine.data.coverage.NearestNeighbourCoverage;
import nz.org.riskscape.engine.function.ExpensiveResource;
import nz.org.riskscape.engine.geo.NearestNeighbourIndex;
import nz.org.riskscape.engine.lookup.LookupTable;
import nz.org.riskscape.engine.problem.GeneralProblems;
import nz.org.riskscape.engine.relation.Relation;
import nz.org.riskscape.engine.types.Referenced;
import nz.org.riskscape.engine.types.Struct;
import nz.org.riskscape.engine.types.Types;
import nz.org.riskscape.hdf5.H5CompoundMember;
import nz.org.riskscape.hdf5.H5Dataset;
import nz.org.riskscape.hdf5.H5DatasetPath;
import nz.org.riskscape.hdf5.H5File;
import nz.org.riskscape.hdf5.relation.H5Relation;
import nz.org.riskscape.hdf5.types.H5CompoundType;
import nz.org.riskscape.oq.events.EventMetadataLookup;
import nz.org.riskscape.oq.gmf.GmfDataLookupBySite;
import nz.org.riskscape.oq.gmf.SitecolGmfDataIndex;
import nz.org.riskscape.oq.sitecol.SitecolRelation;
import nz.org.riskscape.problem.ResultOrProblems;

public class OpenQuakeResolver extends BaseBookmarkResolver<OpenQuakeResolver.Parameters> {

  private static final int CURSOR_PAGE_SIZE_MB = 4;

  public enum DatasetMode {
    /**
     * Turn the sitecol dataset into a coverage that we can index with geometry and find the
     * corresponding site
     */
    COVERAGE_SITE_IDS(TypedCoverage.class),

    /**
     * Return the sitecol dataset site information in relation form
     */
    RELATION_SITE_IDS(Relation.class),

    /**
     * For a given site, get all the event-IDs and GMV data associated with it
     */
    LOOKUP_GMV_BY_SITE(LookupTable.class),

    /**
     * For a given event-ID, get the rupture metadata associated with it (occurrence rate, etc)
     */
    LOOKUP_EVENT_METADATA(LookupTable.class);

    private final Class<?> dataType;

    DatasetMode(Class<?> dataType) {
      this.dataType = dataType;
    }

    public Class<?> getDataType() {
      return dataType;
    }
  }

  public static class Parameters extends BookmarkParameters {

    @ParameterField
    public String sitecolDataset = "/sitecol";

    @ParameterField
    public String gmfIndicesDatasetPath = "/gmf_data/indices";

    @ParameterField
    public String gmfDataDatasetPath = "/gmf_data/data";

    @ParameterField
    public String eventsDatasetPath = "/events";

    @ParameterField
    public String rupturesDatasetPath = "/ruptures";

    @ParameterField
    public DatasetMode mode = DatasetMode.COVERAGE_SITE_IDS;

    @ParameterField(defaultValue = "EPSG:4326")
    public CoordinateReferenceSystem sourceCrs;

    public H5File h5File;

    @ParameterField
    public double maxSiteDistanceMetres = 10000;

    public H5File getH5File() {
      if (h5File == null) {
        h5File = new H5File(validatedFile.get().toPath());
      }
      // TODO when to close? on finalize?
      return h5File;
    }

    @Override
    public Class<?> getDataType() {
      return mode.dataType;
    }

    public H5Dataset openSitecolDataset() {
      return getH5File().openDataset(H5DatasetPath.parse(sitecolDataset));
    }

    public H5Dataset openGmfIndicesDataset() {
      return getH5File().openDataset(H5DatasetPath.parse(gmfIndicesDatasetPath));
    }

    public H5Dataset openGmfDataDataset() {
      return getH5File().openDataset(H5DatasetPath.parse(gmfDataDatasetPath));
    }

    public H5Dataset openEventsDataset() {
      return getH5File().openDataset(H5DatasetPath.parse(eventsDatasetPath));
    }

    public H5Dataset openRupturesDataset() {
      return getH5File().openDataset(H5DatasetPath.parse(rupturesDatasetPath));
    }
  }

  @Getter
  private final Map<String, String> extensionsToFormats = ImmutableMap.of("hdf5", "openquake");

  public OpenQuakeResolver(Engine engine) {
    super(engine);
  }

  @Override
  protected ResultOrProblems build(Parameters parameters) {
    switch (parameters.mode) {
    case COVERAGE_SITE_IDS:
      return buildCoverageAllEvents(parameters);
    case RELATION_SITE_IDS:
      return buildRelationSiteIds(parameters);
    case LOOKUP_GMV_BY_SITE:
      return buildLookupGmvBySite(parameters);
    case LOOKUP_EVENT_METADATA:
      return buildLookupEventMetadata(parameters);

    default:
      throw new UnsupportedOperationException(parameters.mode.toString());
    }

  }

  private ResultOrProblems<LookupTable> buildLookupEventMetadata(Parameters parameters) {
    H5Dataset eventsDataset = parameters.openEventsDataset();
    H5Dataset rupturesDataset = parameters.openRupturesDataset();
    return ResultOrProblems.of(new EventMetadataLookup(eventsDataset, rupturesDataset));
  }

  private ResultOrProblems<LookupTable> buildLookupGmvBySite(Parameters parameters) {
    SitecolGmfDataIndex gmfDataIndex = new SitecolGmfDataIndex(parameters.openGmfIndicesDataset());
    H5Dataset gmfDataset = parameters.openGmfDataDataset();
    return ResultOrProblems.of(new GmfDataLookupBySite(gmfDataIndex, gmfDataset));
  }

  private ResultOrProblems<Relation> buildRelationSiteIds(Parameters parameters) {
    return ResultOrProblems.of(new SitecolRelation(
        parameters.openSitecolDataset(),
        parameters.sourceCrs,
        parameters.getRealizationContext().getProject().getSridSet()
    ));
  }

  protected ResultOrProblems<NearestNeighbourCoverage> buildCoverageAllEvents(Parameters parameters) {

    if (parameters.maxSiteDistanceMetres <= 0) {
      return ResultOrProblems.failed(GeneralProblems.get().
          valueOutOfRange("max-site-distance-metres", parameters.maxSiteDistanceMetres, Range.greaterThan(0D))
      );
    }

    Struct siteType = Struct.of(
        "point", Referenced.of(Types.POINT, parameters.sourceCrs),
        "sid", Types.INTEGER
    );

    CoordinateReferenceSystem crs = parameters.sourceCrs;
    GeometryFactory gm = parameters.getProject().getSridSet().getGeometryFactory(crs);

    H5Dataset siteDataset = parameters.openSitecolDataset();
    H5CompoundType siteDataType = (H5CompoundType) siteDataset.getDataType();
    H5CompoundMember lonMember = siteDataType.findMember("lon");
    H5CompoundMember latMember = siteDataType.findMember("lat");
    AtomicLong indexGenerator = new AtomicLong(0);
    H5Relation siteRelation = new H5Relation(siteDataset, siteType,
        cursor -> {
          Number lon = (Number)cursor.peek(lonMember);
          Number lat = (Number)cursor.peek(latMember);
          return Tuple.ofValues(siteType,
              gm.createPoint(new Coordinate(lat.floatValue(), lon.floatValue())),
              indexGenerator.getAndIncrement()
          );
        }, CURSOR_PAGE_SIZE_MB);


    NearestNeighbourIndex index = NearestNeighbourIndex.metricMaxDistance(
      siteType.getMember("point").get(),
      parameters.getProject().getSridSet(),
      crs,
      parameters.maxSiteDistanceMetres
    );

    ExpensiveResource<NearestNeighbourIndex> resource =
        new ExpensiveResource<NearestNeighbourIndex>(
            parameters.getExecutionContext().getProblemSink(),
            "sitecol-" + parameters.getLocation().toString(), () -> {
              siteRelation.iterator().forEachRemaining(t -> index.insert(t));
              return index;
            }
        );

    return ResultOrProblems.of(
      new NearestNeighbourCoverage(resource, siteType, crs, parameters.getProject().getSridSet())
    );
  }
}
