/*
 * 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.defaults.function;

import static nz.org.riskscape.engine.Matchers.*;
import static org.hamcrest.Matchers.*;
import static org.junit.Assert.*;
import static org.mockito.Mockito.*;

import org.hamcrest.Matcher;
import org.junit.Before;
import org.junit.Test;
import org.locationtech.jts.geom.Point;

import nz.org.riskscape.engine.ArgsProblems;
import nz.org.riskscape.engine.Tuple;
import nz.org.riskscape.engine.coverage.TypedCoverage;
import nz.org.riskscape.engine.data.coverage.MappedCoverage;
import nz.org.riskscape.engine.function.FunctionArgument;
import nz.org.riskscape.engine.function.StringFunctions;
import nz.org.riskscape.engine.function.lang.IfNull;
import nz.org.riskscape.engine.gt.GeometryHelper;
import nz.org.riskscape.engine.problem.GeneralProblems;
import nz.org.riskscape.engine.rl.BaseExpressionRealizerTest;
import nz.org.riskscape.engine.rl.DefaultOperators;
import nz.org.riskscape.engine.rl.MathsFunctions;
import nz.org.riskscape.engine.types.CoverageType;
import nz.org.riskscape.engine.types.Nullable;
import nz.org.riskscape.engine.types.Struct;
import nz.org.riskscape.engine.types.Type;
import nz.org.riskscape.engine.types.Types;
import nz.org.riskscape.problem.Problem;
import nz.org.riskscape.rl.ast.ExpressionProblems;

public class MapCoverageTest extends BaseExpressionRealizerTest {

  GeometryHelper nztmHelper = GeometryHelper.nzTransverseMercatorWithOriginShift(project.getSridSet());
  Point point = nztmHelper.point(1, 1);
  Point pointOutOfBounds = nztmHelper.point(2, 2);
  TypedCoverage mockCoverage;
  Tuple inputTuple = Tuple.EMPTY_TUPLE;
  MappedCoverage coverageResult;
  Type coverageType = new CoverageType(Types.FLOATING);

  @Before
  public void setup() {
    project.getFunctionSet().add(new MapCoverage().identified("map_coverage"));
    project.getFunctionSet().insertFirst(DefaultOperators.INSTANCE);
    project.getFunctionSet().addAll(MathsFunctions.FUNCTIONS);
    project.getFunctionSet().addAll(StringFunctions.FUNCTIONS);
    project.getFunctionSet().add(new IfNull().identified("if_null"));

    mockCoverage = mock(TypedCoverage.class, "foo");
    when(mockCoverage.getType()).thenReturn(Types.FLOATING);
    when(mockCoverage.getCoordinateReferenceSystem()).thenReturn(nztmHelper.getCrs());
    when(mockCoverage.evaluate(point)).thenReturn(10.0D);
    when(mockCoverage.evaluate(pointOutOfBounds)).thenReturn(null);

    inputStruct = Struct.EMPTY_STRUCT.add("foo", coverageType);
    inputTuple = Tuple.ofValues(inputStruct.asStruct(), mockCoverage);
  }

  @Test
  public void canPerformSimpleMapping() throws Exception {
    evaluate("map_coverage(foo, x -> x + 1)");
    assertSuccessful();

    assertThat(coverageResult.getType(), is(Types.FLOATING));
    assertThat(coverageResult.evaluate(point), is(11.0D));
    assertNull(coverageResult.evaluate(pointOutOfBounds));
  }

  @Test
  public void canPerformComplexMapping() throws Exception {
    evaluate("map_coverage(foo, val -> { raw: val, mapped: str(square_root(val - 1)) })");
    assertSuccessful();

    Struct expected = Struct.of("raw", Types.FLOATING, "mapped", Types.TEXT);
    assertThat(coverageResult.getType(), is(expected));
    assertThat(coverageResult.evaluate(point), is(Tuple.ofValues(expected, 10.0D, "3.0")));
    assertNull(coverageResult.evaluate(pointOutOfBounds));
  }

  @Test
  public void canPreserveNullabilityOfContainedType() throws Exception {
    inputStruct = Struct.EMPTY_STRUCT.add("foo", new CoverageType(Nullable.FLOATING));
    inputTuple = Tuple.ofValues(inputStruct.asStruct(), mockCoverage);
    when(mockCoverage.getType()).thenReturn(Nullable.FLOATING);

    // sanity-check the returned coverage itself is not nullable
    realize(inputStruct, parse("map_coverage(foo, val -> round(val * 2))"));
    assertThat(realized.getResultType(), is(new CoverageType(Nullable.INTEGER)));

    // expression is not null-safe, so the coverage should still hold a nullable type
    evaluate("map_coverage(foo, val -> round(val * 2))");
    assertSuccessful();
    assertThat(coverageResult.getType(), is(Nullable.INTEGER));
    assertThat(coverageResult.evaluate(point), is(20L));
    assertNull(coverageResult.evaluate(pointOutOfBounds));
  }

  @Test
  public void canDiscardNullabilityOfContainedType() throws Exception {
    inputStruct = Struct.EMPTY_STRUCT.add("foo", new CoverageType(Nullable.FLOATING));
    inputTuple = Tuple.ofValues(inputStruct.asStruct(), mockCoverage);
    when(mockCoverage.getType()).thenReturn(Nullable.FLOATING);

    // sanity-check the returned coverage itself is not nullable
    realize(inputStruct, parse("map_coverage(foo, val -> str(if_null(val, -1.0)))"));
    assertThat(realized.getResultType(), is(new CoverageType(Types.TEXT)));

    // null-safe expression so the coverage should contain a non-nullable type
    evaluate("map_coverage(foo, val -> str(if_null(val, -1.0)))");
    assertSuccessful();
    assertThat(coverageResult.getType(), is(Types.TEXT));
    assertThat(coverageResult.evaluate(point), is("10.0"));
    // when the coverage returns null it should still get mapped
    assertThat(coverageResult.evaluate(pointOutOfBounds), is("-1.0"));
  }

  @Test
  public void canPreserveNullabilityOfCoverage() throws Exception {
    inputStruct = Struct.EMPTY_STRUCT.add("foo", Nullable.of(coverageType));
    inputTuple = Tuple.ofValues(inputStruct.asStruct(), mockCoverage);

    // the expression should preserve the nullable nature of the coverage itself
    realize(inputStruct, parse("map_coverage(foo, val -> str(if_null(val, -1.0)))"));
    assertThat(realized.getResultType(), is(Nullable.of(new CoverageType(Types.TEXT))));

    evaluate("map_coverage(foo, val -> str(if_null(val, -1.0)))");
    assertSuccessful();

    // the coverage itself yields a non-nullable type still
    assertThat(coverageResult.getType(), is(Types.TEXT));
    assertThat(coverageResult.evaluate(point), is("10.0"));
    // note that underlying contained type is not nullable, so nulls aren't mapped
    // (even though our lambda expression would handle nulls)
    assertNull(coverageResult.evaluate(pointOutOfBounds));

    // passing in a null coverage should produce a null result
    inputTuple = Tuple.of(inputStruct.asStruct(), "foo", null);
    evaluate("map_coverage(foo, val -> str(if_null(val, -1.0)))");
    assertTrue(realizationProblems.isEmpty());
    assertNull(evaluated);
  }

  @Test
  public void cannotMapWithoutCoverage() throws Exception {
    evaluate("map_coverage()");
    assertFailed(hasAncestorProblem(is(
        GeneralProblems.get().required(expectedArg(0)
    ))));

    evaluate("map_coverage('foo', x -> x + 1)");
    assertFailed(hasAncestorProblem(is(
        ArgsProblems.mismatch(expectedArg(0), Types.TEXT)
    )));
  }

  @Test
  public void cannotMapWithoutLambdaExpression() throws Exception {
    evaluate("map_coverage(foo)");
    assertFailed(hasAncestorProblem(is(
        GeneralProblems.get().required(expectedArg(1))
    )));

    evaluate("map_coverage(foo, 'bar')");
    assertFailed(hasAncestorProblem(is(
        ArgsProblems.mismatch(expectedArg(1), Types.TEXT)
    )));
  }

  @Test
  public void canReferenceAttributesOutsideLambdaScope() throws Exception {
    inputStruct = ((Struct) inputStruct).add("baz", Types.FLOATING);
    inputTuple = Tuple.ofValues(inputStruct.asStruct(), mockCoverage, 50.0D);

    evaluate("map_coverage(foo, x -> x + baz)");
    assertSuccessful();

    assertThat(coverageResult.getType(), is(Types.FLOATING));
    assertThat(coverageResult.evaluate(point), is(60.0D));
    assertNull(coverageResult.evaluate(pointOutOfBounds));
  }

  @Test
  public void singleArgsLambdaExpressionRequired() throws Exception {
    String badLambda = "(x, y) -> x * y + 1";
    evaluate("map_coverage(foo, " + badLambda + ")");
    assertFailed(hasAncestorProblem(
        is(ExpressionProblems.get().lambdaArityError(parse(badLambda), 2, 1)
    )));
  }

  private void assertFailed(Matcher<Problem> problemMatcher) {
    assertThat(realizationProblems, contains(problemMatcher));
  }

  private void assertSuccessful() {
    if (!realizationProblems.isEmpty()) {
      fail(Problem.debugString(realizationProblems));
    }
    assertThat(evaluated, instanceOf(MappedCoverage.class));
    coverageResult = (MappedCoverage) evaluated;
  }

  private void evaluate(String expr) {
    evaluate(expr, inputTuple);
  }

  private FunctionArgument expectedArg(int index) {
    return project.getFunctionSet().get("map_coverage").getArguments().get(index);
  }
}
