/*
 * 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 static org.junit.Assert.*;

import java.util.ArrayList;
import java.util.HashMap;
import java.util.HashSet;
import java.util.List;
import java.util.Set;

import org.junit.After;
import org.junit.Before;
import org.junit.Test;
import org.geotools.api.referencing.crs.CoordinateReferenceSystem;

import nz.org.riskscape.engine.CrsHelper;
import nz.org.riskscape.engine.SRIDSet;
import nz.org.riskscape.hdf5.H5Dataset;
import nz.org.riskscape.hdf5.H5File;
import nz.org.riskscape.oq.gmf.GmfData;
import nz.org.riskscape.oq.gmf.GmfDataBySiteIdList;
import nz.org.riskscape.oq.gmf.SitecolGmfDataIndex;
import nz.org.riskscape.oq.sitecol.SitecolIterator;

public class SitecolIndexTest extends BaseHdf5Test implements CrsHelper {

  SRIDSet sridSet = new SRIDSet();
  CoordinateReferenceSystem crs = longLat();
  private H5File h5file;
  private H5Dataset indexDataset;
  private H5Dataset gmfDataset;
  private H5Dataset sitecolDataset;
  private SitecolGmfDataIndex index;

  @Before
  public void openDatasets() {
    h5file = probabilisticFile();
    indexDataset = h5file.openDataset("/gmf_data", "indices");
    gmfDataset = h5file.openDataset("/gmf_data", "data");
    sitecolDataset = h5file.openDataset("/", "sitecol");
    index = new SitecolGmfDataIndex(indexDataset);
  }

  @After
  public void closeDatasets() {
    sitecolDataset.close();
    gmfDataset.close();
    indexDataset.close();
    h5file.close();
  }

  @Test
  public void canLookUpAnEmptySite() {
    // there's a ton of empties in the test dataset, I've picked one at ... randomish
    GmfDataBySiteIdList list = new GmfDataBySiteIdList(100, index, gmfDataset);
    assertEquals(0, list.size());
  }

  @Test
  public void canLookupASiteWithASingleEvent() throws Exception {
    GmfDataBySiteIdList list = new GmfDataBySiteIdList(14686, index, gmfDataset);
    assertEquals(1, list.size());
    GmfData gmf = list.get(0);

    // values taken from h5dumping the file
    assertEquals(14686L, gmf.siteId);
    assertEquals(57L, gmf.eventId);
    assertEquals(0.0161901, gmf.gmv[0], 0.000001);
  }


  @Test
  public void canLookupASiteWithAFewContiguousEvents() throws Exception {
    GmfDataBySiteIdList list = new GmfDataBySiteIdList(48864, index, gmfDataset);
    assertEquals(3, list.size());

    // values taken from h5dumping the file
    float[] expectedShaking = new float[] {0.0504928F, 0.0349537F, 0.0339037F};
    long[] expectedEventIds = new long[] {60L, 61L, 62L};

    int idx = 0;
    for (GmfData gmfData : list) {
      assertEquals(Integer.toString(idx), 48864L, gmfData.siteId);
      assertEquals(Integer.toString(idx), expectedShaking[idx], gmfData.gmv[0], 0.000001);
      assertEquals(Integer.toString(idx), expectedEventIds[idx], gmfData.eventId);

      idx++;
    }
  }

  @Test
  public void canLookupASiteWithManyEvents() throws Exception {
    long siteId = 14868;
    long[][] indices = index.lookupIndices(siteId);
    // three distinct regions in the gmf data dataset that we index through
    assertEquals(3, indices[0].length);
    assertEquals(3, indices[1].length);
    GmfDataBySiteIdList list = new GmfDataBySiteIdList(siteId, indices, gmfDataset);

    list.get(3);

    // values taken from h5dumping
    long[] expectedEventIds = new long[] {60L, 61L, 62L, 57L, 63L, 64L, 65L, 66L, 67L};
    float[] expectedShaking = new float[] {0.0495685F, 0.0343138F, 0.033283F, 0.0190081F,
        0.0371584F, 0.045341F, 0.0335395F, 0.0351488F, 0.0206659F};

    assertEquals(9, list.size());

    int idx = 0;
    // we'll compare these to some out of order selection
    List<GmfData> collected = new ArrayList<>();
    for (GmfData gmfData : list) {
      assertEquals(Integer.toString(idx), siteId, gmfData.siteId);
      assertEquals(Integer.toString(idx), expectedShaking[idx], gmfData.gmv[0], 0.000001);
      assertEquals(Integer.toString(idx), expectedEventIds[idx], gmfData.eventId);

      collected.add(gmfData);
      idx++;
    }

    int[][] shuffledIndexes = new int[][] {
      {8, 7, 6, 5, 4, 3, 2, 1, 0}, // reverse
      {8, 0, 7, 1, 6, 2, 5, 3, 4}, // back and forth
      {0, 2, 4, 6, 8, 1, 3, 5, 7}, // evens odds
      {0, 2, 4, 6, 8, 7, 5, 3, 1}, // evens, then reverse odds
      {7, 2, 3, 6, 4, 1, 8, 5, 0}, // random1
      {3, 1, 8, 0, 4, 6, 5, 2, 7}, // random2
    };

    // now do some random access through the list and check it matches sequential access
    for (int orderingIndex = 0; orderingIndex < shuffledIndexes.length; orderingIndex++) {
      int[] ordering = shuffledIndexes[orderingIndex];
      for (int orderIndex = 0; orderIndex < ordering.length; orderIndex++) {
        int mappedIdx = ordering[orderIndex];
        assertEquals(collected.get(mappedIdx), list.get(mappedIdx));
      }
    }
  }

  private HashMap<Long, GmfDataBySiteIdList> readAllSiteData() {
    HashMap<Long, GmfDataBySiteIdList> siteData = new HashMap<>();
    SitecolIterator iterator = new SitecolIterator(sitecolDataset);

    while (iterator.hasNext()) {
      long siteId = iterator.next().siteId;
      GmfDataBySiteIdList list = new GmfDataBySiteIdList(siteId, index, gmfDataset);
      siteData.put(siteId, list);
    }
    return siteData;
  }

  private void assertSiteDataValid(HashMap<Long, GmfDataBySiteIdList> siteData) {
    for (long siteId : siteData.keySet()) {
      GmfDataBySiteIdList list = siteData.get(siteId);
      assertTrue(Long.toString(siteId), list.size() < 69);

      if (list.size() > 0) {

        Set<Long> eventIds = new HashSet<>();
        for (GmfData gmfData : list) {
          assertEquals(gmfData.siteId, siteId);
          // events seem to be sorted in ascending order per site
          eventIds.add(gmfData.eventId);
          // sanity check
          assertTrue(gmfData.eventId >= 0);
          assertTrue(gmfData.eventId < 68);
        }

        // check we got the correct number of events and each event ID is unique
        assertEquals(list.size(), eventIds.size());
        HashSet<Long> uniqueEventIds = new HashSet<>(eventIds);
        assertEquals(list.size(), uniqueEventIds.size());
      }
    }
  }

  @Test
  public void canLookUpAllTheIndexesInOrder() {
    HashMap<Long, GmfDataBySiteIdList> siteData = readAllSiteData();
    assertSiteDataValid(siteData);
  }

  @Test
  public void canReadSiteIndicesInParallel() throws Exception {
    // check we can read the site indices data from different threads without corruption
    List<HashMap<Long, GmfDataBySiteIdList>> results = new ArrayList<>();

    // run two threads in parallel that both try to read the same site data
    Thread t1 = new Thread(() -> {
      HashMap<Long, GmfDataBySiteIdList> data = readAllSiteData();
      synchronized (results) {
        results.add(data);
      }
    });
    Thread t2 = new Thread(() -> {
      HashMap<Long, GmfDataBySiteIdList> data = readAllSiteData();
      synchronized (results) {
        results.add(data);
      }
    });

    t1.start();
    t2.start();

    try {
      t1.join();
      t2.join();
    } catch (InterruptedException e) {
    }

    HashMap<Long, GmfDataBySiteIdList> siteData1 = results.get(0);
    HashMap<Long, GmfDataBySiteIdList> siteData2 = results.get(0);
    assertSiteDataValid(siteData1);
    assertSiteDataValid(siteData2);
    assertEquals(siteData1, siteData2);
  }

  @Test
  public void canAlsoIndexInToScenarioHdf5() throws Exception {
    h5file.close();
    h5file = scenarioFile();
    indexDataset = h5file.openDataset("/gmf_data", "indices");
    index = new SitecolGmfDataIndex(indexDataset);

    // just so happens that the locations are predictable, but we index anyway, because why not
    assertStartStop(0, 0, 10);
    assertStartStop(20, 200, 210);
  }

  private void assertStartStop(long siteId, long start, long stop) {
    long[][] locations = index.lookupIndices(siteId);

    assertEquals(2, locations.length);
    assertEquals(1, locations[0].length);
    assertEquals(1, locations[1].length);

    assertEquals(start, locations[0][0]);
    assertEquals(stop, locations[1][0]);

  }

}
