import "utilsLib" as utilsLib;
import "numberLib" as numberLib;
import "journalTypes" as journalTypes;

def ingestionDose($customUnits):
  . as $ingestion |
  journalTypes::ensureIngestion |

  if .customUnitId != null then
    ($customUnits | map(select(.id == $ingestion.customUnitId))[0]) as $customUnit |
    .dose * $customUnit.dose | . as $dose | 
    $dose
  else 
    .dose
  end;

def addStandardDeviations($expectationX; $standardDeviationX; $expectationY; $standardDeviationY):
  (pow($standardDeviationX; 2) + pow($expectationX; 2)) as $sumX |
  (pow($standardDeviationY; 2) + pow($expectationY; 2)) as $sumY |
  (pow($expectationX; 2) * pow($expectationY; 2)) as $expectations |
  ($sumX * $sumY - $expectations) as $productVariance |
  if $productVariance > 0.0000001 then
    $productVariance | sqrt | numberLib::round(2)
  else 
    null
  end;

def ingestionStandardDeviation($customUnits):
  . as $ingestion |
  journalTypes::ensureIngestion |

  if .customUnitId != null then
    ($customUnits | map(select(.id == $ingestion.customUnitId))[0]) as $customUnit |

    ($ingestion.dose // 0) as $expectationX |
    ($ingestion.estimatedDoseStandardDeviation // 0) as $standardDeviationX |
    ($customUnit.dose // 0) as $expectationY |
    ($customUnit.estimatedDoseStandardDeviation // 0) as $standardDeviationY |

    addStandardDeviations($expectationX; $standardDeviationX; $expectationY; $standardDeviationY)
  else 
    .estimatedDoseStandardDeviation
  end;

def ingestionUnit($customUnits):
  . as $ingestion |
  journalTypes::ensureIngestion |
  if .customUnitId != null then
    ($customUnits | map(select(.id == $ingestion.customUnitId))[0]) as $customUnit |
    $customUnit.originalUnit
  else 
    .units
  end;

def ingestionConsumerName:
  . as $ingestion |
  journalTypes::ensureIngestion |
  $ingestion.consumerName // "default";

def filterIngestions($substanceFilter; $consumerFilter):
  . as $ingestions |
  ($substanceFilter // []) as $substanceFilter |
  (($substanceFilter | length) > 0) as $shouldFilterBySubstance |
  ($consumerFilter // []) as $consumerFilter |
  (($consumerFilter | length) > 0) as $shouldFilterByConsumer |

  journalTypes::ensureIngestions |

  if $shouldFilterBySubstance or $shouldFilterByConsumer then 
    [
      $ingestions[] as $ingestion |
        if 
          (if $shouldFilterBySubstance then ([$substanceFilter[] | . == $ingestion.substanceName] | any) else true end
          and
          if $shouldFilterByConsumer then ([$consumerFilter[] | . == ($ingestion | ingestionConsumerName)] | any) else true end)
        then $ingestion
        else null end
    ] | map(select(. != null))
  else $ingestions end;

def ingestionsSubstanceNames:
  . as $ingestions |
  journalTypes::ensureIngestions |
  [$ingestions[].substanceName] | utilsLib::orderedUnique;

def ingestionsConsumerNames:
  . as $ingestions |
  journalTypes::ensureIngestions |
  [$ingestions[] | ingestionConsumerName] | utilsLib::orderedUnique;

def ingestionsByConsumer:
  . as $ingestions |
  journalTypes::ensureIngestions |
  ($ingestions | ingestionsConsumerNames) as $consumerNames |
  [
    $consumerNames[] as $consumerName |
      {
        key: $consumerName,
        value: $ingestions | map(select((. | ingestionConsumerName) == $consumerName)),  
      }
  ] | from_entries;

def defaultIngestionInfo:
  {
    unit: "",
    dose: null,
    isUnknown: false,
    isEstimate: false,
    # TODO
    standardDeviation: null
  };

def ingestionInfo($customUnits):
  . as $ingestion |
  .substanceName as $name |
  .administrationRoute as $administrationRoute |
  . | ingestionStandardDeviation($customUnits) as $standardDeviation |
  . | ingestionDose($customUnits) as $dose |
  . | ingestionUnit($customUnits) as $unit |
  . | ingestionConsumerName as $consumerName |

  defaultIngestionInfo as $ingestionInfo |
  $ingestionInfo |

  if ($dose == null) then .isUnknown |= true end |

  if ($standardDeviation != null) then
    .isEstimate |= true
  end |

  .unit |= $unit |
  .dose |= $dose |
  .standardDeviation |= $standardDeviation;

def addIngestionInfo($rhs):
  . as $lhs | 

  if (.unit != "" and $rhs.unit != "") then
    utilsLib::assert(.unit == $rhs.unit; "units mismatch, todo add reasonable conversions")
  else 
    .unit |= $rhs.unit
  end |


  if ($rhs.dose == null) then .isUnknown |= true end |

  if ($rhs.standardDeviation != null) then
    .isEstimate |= true 
  end |

  .dose += $rhs.dose |
  .standardDeviation += $rhs.standardDeviation;

def addIngestionInfos:
  . as $ingestionInfos |
  reduce $ingestionInfos[] as $ingestionInfo (defaultIngestionInfo;
    . | addIngestionInfo($ingestionInfo)
  );

def experienceStats($customUnits):
  . as $experience |
  journalTypes::ensureExperience |
  $experience.ingestions as $ingestions |
  (reduce $ingestions[] as $ingestion ({}; . as $stats |
    $ingestion |
    .substanceName as $name |
    .administrationRoute as $administrationRoute |
    (. | ingestionConsumerName) as $consumerName |

    $stats |
      .[$consumerName].[$ingestion.substanceName].[$administrationRoute] as $ingestionStats |

    if $ingestionStats == null then
      ($ingestion | ingestionInfo($customUnits))
    else
      $ingestionStats | addIngestionInfo($ingestion | ingestionInfo($customUnits))
    end | . as $ingestionStats |

    $stats |
      .[$consumerName].[$ingestion.substanceName].[$administrationRoute] |= $ingestionStats
  ));

def statsCalculateCombinedDose($substanceName; $consumerName):
  . as $stats |
  utilsLib::assert($stats | has($consumerName); "consumer not found in stats") |
  utilsLib::assert($stats[$consumerName] | has($substanceName); 
    "substance for consumer not found in stats"
  ) |

  .[$consumerName].[$substanceName] | values | addIngestionInfos;

def experiencesWithExtraData($customUnits):
  . as $experiences |
  journalTypes::ensureExperiences |
  (reduce $experiences[] as $experience ([]; 
    ($experience | experienceStats($customUnits)) as $stats |
    . += [{
      $stats,
      $experience
    }]
  ));

def filterSortExperiences($customUnits; $substanceFilter; $consumerFilter; $sortMethod; $sortOptions):
  . as $experiences |
  journalTypes::ensureExperiences |
  ($consumerFilter // ["default"]) as $consumerFilter |
  $sortOptions |
    # If filtering results by substances but no sortOptions.substanceName is defined, use the first one as a default
    (.substanceName |= 
      if 
        $sortOptions.substanceName == null
        and
        (($substanceFilter | length) >= 1)
      then
        $substanceFilter[0]
      else
        null
      end
  ) | . as $sortOptions |
  $sortOptions | (.ingestionMethod |= ($sortOptions.ingestionMethod // "ORAL")) | . as $sortOptions |
  $sortOptions | (.consumerName |= ($sortOptions.consumerName // "default")) | . as $sortOptions |

  def sortFilter:
    def oldToNewSort: .experience.sortDate;
    def highestCombinedDoseSort: .stats.[$sortOptions.consumerName].[$sortOptions.substanceName] | [to_entries[] | .value.dose] | add;
    def highestMethodDoseSort: .stats.[$sortOptions.consumerName].[$sortOptions.substanceName].[$sortOptions.ingestionMethod // "ORAL"].dose;

    def filterBySubstanceFilter:
      . as $experiencesData |
      [
        $experiencesData[] as $experienceData |
        ($experienceData.experience.ingestions) as $ingestions |
          if ($experienceData.experience.ingestions |
            any(
              . as $ingestion |
              $substanceFilter | 
              any(index($ingestion.substanceName))
            )
          ) then $experienceData else null end
      ] | map(select(. != null));

    def filterByConsumerFilter:
      . as $experiencesData |
      (reduce $experiencesData[] as $experienceData ([]; 
        ($experienceData.experience.ingestions) as $ingestions |
        . += 
            if ($experienceData.experience.ingestions |
              any(
                (. | ingestionConsumerName) as $consumerName |
                $consumerFilter |
                  any(index($consumerName))
              )
            ) then [$experienceData] else [] end
      ));

    def filterBySubstanceAndConsumer:
      . as $experiencesData |
      (reduce $experiencesData[] as $experienceData ([]; 
        ($experienceData.experience.ingestions) as $ingestions |
        . += 
            if 
              ($experienceData.experience.ingestions |
                any(
                    .substanceName == $sortOptions.substanceName
                      and
                    (. | ingestionConsumerName) == $sortOptions.consumerName
                  ))
            then [$experienceData] else [] end
      ));

    def filterByIngestionMethodForSubstance:
      . as $experiencesData | filterBySubstanceAndConsumer as $experiencesData |
      (reduce $experiencesData[] as $experienceData ([]; 
        ($experienceData.experience.ingestions) as $ingestions |
        . += if ($experienceData.experience.ingestions | any(.substanceName == $sortOptions.substanceName and .administrationRoute == $sortOptions.ingestionMethod))
          then [$experienceData] else [] end
      ));

    # speeds up by excluding everything not containing substances & consumers not in filters, wouldn't show any data anyway
    if $substanceFilter != null then filterBySubstanceFilter end |
    filterByConsumerFilter |

    if 
      $sortMethod == "old-to-new" or $sortMethod == null
    then 
      sort_by(oldToNewSort)
    elif 
      $sortMethod == "highest-combined-dose"
    then
      utilsLib::assert($sortOptions.substanceName != null; "substanceName not provided as sortOption") |
      filterBySubstanceAndConsumer | 
        sort_by(highestCombinedDoseSort, oldToNewSort) |
          reverse
    elif 
      $sortMethod == "highest-dose-for-method"
    then
      utilsLib::assert($sortOptions.substanceName != null; "substanceName not provided as sortOption") |
      filterByIngestionMethodForSubstance | 
        sort_by(highestMethodDoseSort, oldToNewSort) |
          reverse
    end;

  $experiences |
    experiencesWithExtraData($customUnits) |
    sortFilter;

def experienceByTitle($name):
  utilsLib::assert((. | type) == "array"; "experienceByTitle takes a array of experiences as input") |
  map(select(.title == $name))[0];