/*
 * 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.wizard.model2.analysis;

import static nz.org.riskscape.rl.TokenTypes.*;
import static nz.org.riskscape.wizard.model2.DslHelper.*;
import static nz.org.riskscape.wizard.model2.smp.SamplePhase.*;

import java.util.ArrayList;
import java.util.Arrays;
import java.util.EnumSet;
import java.util.List;
import java.util.Optional;
import java.util.stream.Collectors;

import nz.org.riskscape.engine.function.IdentifiedFunction;
import nz.org.riskscape.engine.rl.FunctionCallPrototype;
import nz.org.riskscape.engine.types.RSList;
import nz.org.riskscape.engine.types.Struct;
import nz.org.riskscape.engine.types.Type;
import nz.org.riskscape.rl.ast.PropertyAccess;
import nz.org.riskscape.wizard.AskerHints;
import nz.org.riskscape.wizard.QuestionSet;
import nz.org.riskscape.wizard.Survey;
import nz.org.riskscape.wizard.bld.IncrementalBuildState;
import nz.org.riskscape.wizard.bld.PipelineChange;
import nz.org.riskscape.wizard.bld.dsl.IncompletePipelineChange;
import nz.org.riskscape.wizard.bld.dsl.PipelineChangeInput;
import nz.org.riskscape.wizard.bld.loc.ChangeLocation;
import nz.org.riskscape.wizard.model2.input.InputDataPhase;
import nz.org.riskscape.wizard.model2.smp.SamplePhase;
import nz.org.riskscape.wizard.survey2.BasePhase;
import nz.org.riskscape.wizard.survey2.Choices;
import nz.org.riskscape.wizard.survey2.DefaultQuestionSet2;

public class AnalysisPhase extends BasePhase {

  public AnalysisPhase(Survey survey) {
    super(survey);
  }

  public static final String CONSEQUENCE = "consequence";

  private static final String QUESTION_SET_ID = "analysis";

  @Override
  public QuestionSet buildQuestionSet(IncrementalBuildState buildState) {

    DefaultQuestionSet2 qs = new DefaultQuestionSet2(QUESTION_SET_ID, this);

    qs.setDefaultLocation(InputDataPhase.MAIN_BRANCH);

    qs.addQuestion("map-hazard", Boolean.class)
      .customizeQuestion(q -> q.askWhen(
          // no point asking this for multi-hazard, as we always end up with a struct
          bs -> !InputDataPhase.isMultiHazard(bs) && getHazardType(bs.getInputStruct(q)).isA(Struct.class)))
      .thenNoChange();

    qs.addQuestion("map-hazard-intensity", PropertyAccess.class)
      .askWhenResponseIs("map-hazard", Boolean.TRUE)
      .customizeQuestion(q -> q.withChoices(ibs -> {
        Type hazardType = getHazardType(ibs.getInputStruct(q));
        return Choices.forScope(hazardType.findAllowNull(Struct.class).get(), q);
      }))
      .then((input, hazardAttr) -> PipelineChange.chainStep(
            "select({*, %s: map(%s, hv -> hv.%s)})",
            InputDataPhase.HAZARD_ATTRIBUTE,
            InputDataPhase.HAZARD_ATTRIBUTE,
            hazardAttr.toSource()
      ));

    qs.addQuestion("aggregate-hazards", Boolean.class)
      .customizeQuestion(q -> q.askWhen(bs -> {
        // can aggregate a simple list of hazard values, but don't support aggregating struct hazards
        // in the wizard - it gets too tricky, e.g. what is the 'max' when the struct holds multiple values
        Optional<RSList> hazard = getAttributeType(bs.getInputStruct(q), InputDataPhase.HAZARD_ATTRIBUTE, RSList.class);
        return (hazard.isPresent() && !hazard.get().getMemberType().findAllowNull(Struct.class).isPresent());
      }))
      .thenNoChange();

    qs.addQuestion("aggregate-hazards-function", FunctionCallPrototype.class)
      .customizeQuestion(q -> q.withAnnotations(AskerHints.AGGREGATION, ""))
      .askWhenResponseIs("aggregate-hazards", Boolean.TRUE)
      .then((input, aggFunction) -> PipelineChange.chainStep(
        "select({*, %s: %s})",
        InputDataPhase.HAZARD_ATTRIBUTE,
        aggFunction.setFirstArgument(InputDataPhase.HAZARD_ATTRIBUTE).build().toSource()
    ));

    qs
    .addQuestion("function", IdentifiedFunction.class)
    .customizeQuestion(q -> q.withChoices(ibs -> Choices.forFunctions(ibs, CATEGORIES)))
    .then((input, answer) -> {
      // the resource-layer can optionally be passed to the function, if specified
      boolean hasResource = input.getStruct().hasMember(InputDataPhase.RESOURCE_ATTRIBUTE);

      return PipelineChange.chainStep(
          // we can map a list or just apply the expression directly with the map function
          "select({*, %s: map(%s, hv -> %s(%s, hv%s))})",
          CONSEQUENCE,
          InputDataPhase.HAZARD_ATTRIBUTE,
          quoteIdent(answer.getId()),
          InputDataPhase.EXPOSURE_ATTRIBUTE,
          hasResource ? ", " + InputDataPhase.RESOURCE_ATTRIBUTE : ""
      );
    });

    // only scale losses if aggregating the consequence. Scaling when using aggregate-hazard
    // doesn't entirely make sense, as generally when you'd want to scale (such as with a road)
    // the aggregated hazard intensity may not have impacted the exposure uniformly
    // (e.g. max flood depth)
    qs
    .addQuestion("scale-losses", Boolean.class)
        .customizeQuestion(q -> q.askWhen(bs ->
            bs.getInputStruct(q).hasMember(SamplePhase.EXPOSED_RATIO)
            && getAttributeType(bs.getInputStruct(q), CONSEQUENCE, RSList.class).isPresent())
    ).then((input, doScaling) -> {
      if (doScaling) {
        return PipelineChange.chainStep(""
            + "select({*, raw_consequence: consequence,"
            + " consequence: zip(consequence, exposed_ratio, (loss, ratio) -> loss * ratio)})");
      } else {
        // ignore if the user selected the question, then changed their minds
        return PipelineChange.noChange();
      }
    });

    // add a known step name we can save the raw results from
    qs.addHiddenQuestion("raw-results")
      .then(input -> PipelineChange.chainStep("select({*}) as raw_results"));

    qs
    .addQuestion("aggregate-consequences", Boolean.class)
      .customizeQuestion(q -> q.askWhen(bs ->
        getAttributeType(bs.getInputStruct(q), InputDataPhase.HAZARD_ATTRIBUTE, RSList.class).isPresent()
      ))
      .thenNoChange();

    qs.addQuestion("aggregate-consequences-function", FunctionCallPrototype.class)
      .customizeQuestion(q -> q.withAnnotations(AskerHints.AGGREGATION, ""))
      .askWhenResponseIs("aggregate-consequences", Boolean.TRUE)
      .then((input, aggFunction) -> PipelineChange.chainStep(
        "select({*, %s: %s})",
        CONSEQUENCE,
        aggFunction.setFirstArgument(CONSEQUENCE).build().toSource()
    ));

    // all-intersections makes it less obvious how the consequence was determined, especially
    // if we are scaling losses. Saving the raw results lets us 'show our working' (note that
    // it will potentially include duplicate losses for the same asset though)
    qs
    .addQuestion("save-raw-results", Boolean.class)
    .customizeQuestion(q -> q
        .optionalOne()
        .askWhen(bs -> getAttributeType(bs.getInputStruct(q), SamplePhase.SAMPLED_HAZARD, RSList.class).isPresent())
        .atLocation(ChangeLocation.atStep("raw_results"))
    ).then((input, doSave) -> {
      if (doSave) {
        return saveRawResults(input);
      } else {
        return PipelineChange.noChange();
      }
    });

    qs
    .addHiddenQuestion("final")
    .then((input, answer) -> {

      // Get a list of attributes we started with as a working copy for a select expression change
      ArrayList<String> mapAttributes = input.getStruct().getMembers()
          // would have used .getMemberKeys() here, but that returns them sorted
          // which breaks all our existing tests
          .stream().map(Struct.StructMember::getKey)
          .collect(Collectors.toCollection(ArrayList::new));

      // the results may contain an exposed_ratio list, but this is not that useful to users.
      // Convert it to a single total exposed_ratio value instead
      if (getAttributeType(input.getStruct(), EXPOSED_RATIO, RSList.class).isPresent()) {
        // make sure the total sum doesn't exceed 1. This could occur due to
        // precision/reprojection issues, but it'd make the results look a little dodgy
        mapAttributes.set(
            mapAttributes.indexOf(EXPOSED_RATIO),
            String.format("min(sum(%s), 1.0) as %s", EXPOSED_RATIO, EXPOSED_RATIO));
      }

      /*
       * NB We are purposely removing the sampled data here, but not the resulting hazard value - that remains.
       * This data is not particularly useful output (generally speaking) and at best blows out the size of the final
       * output, and at worst can crash the model ( See #1322). If users really want it, they can choose to save the
       * raw analysis results and get it that way.
       */
      mapAttributes.remove(SamplePhase.SAMPLED_HAZARD);

      return PipelineChange.chainStep("select({%s}) as %s", String.join(", ", mapAttributes), LAST_STEP_NAME);
    });

    return qs;
  }

  private static final EnumSet<IdentifiedFunction.Category> CATEGORIES = EnumSet.of(
    IdentifiedFunction.Category.RISK_MODELLING,
    IdentifiedFunction.Category.UNASSIGNED
  );

  public static final String LAST_STEP_NAME = "event_impact_table";

  private IncompletePipelineChange saveRawResults(PipelineChangeInput input) {
    List<String> possibleLists = Arrays.asList(
        SamplePhase.SAMPLED_HAZARD,
        InputDataPhase.HAZARD_ATTRIBUTE,
        "raw_consequence",
        SamplePhase.EXPOSED_RATIO,
        CONSEQUENCE);

    List<String> unnestAttrs = new ArrayList<>();
    Struct inputScope = input.getStruct();

    // unnest any lists, so the user sees each value as a separate row
    for (String attr : possibleLists) {
      if (getAttributeType(inputScope, attr, RSList.class).isPresent()) {
        unnestAttrs.add(attr);
      }
    }

    String unnestStep = unnestAttrs.stream()
        .collect(Collectors.joining(", ", "unnest([", "])"));

    // reorder the attributes so that the hazard_sampled is first, and so its
    // geometry will get used when saving the shapefile
    List<String> attributes = new ArrayList<>(inputScope.getMemberKeys());
    attributes.remove(SamplePhase.SAMPLED_HAZARD);
    attributes.add(0, SamplePhase.SAMPLED_HAZARD);
    String selectStep = String.format("select({%s})", String.join(", ", attributes));

    return PipelineChange.newChain("raw_results -> %s -> %s -> save(name: 'raw-analysis')", unnestStep, selectStep);
  }

  /**
   * @return the type supplied by the hazard layer
   */
  private Type getHazardType(Struct inputType) {
    Type hazardType = getAttributeType(inputType, InputDataPhase.HAZARD_ATTRIBUTE, Type.class).get();
    return hazardType.findAllowNull(RSList.class)
        .map(RSList::getContainedType)
        .orElse(hazardType);
  }
}
