# run as: jq export.json -f tool.jq -Cr --args -- 

include "utils";
include "argsLib";
include "journalUtils";
import "types" as types;
import "typeLib" as typeLib;

def printExperienceStats($stats; $substanceFilter; $consumerFilter; $withTitle):
  . as $experience |

  ($consumerFilter // ["default"]) as $consumerFilter |

  $experience.ingestions | filterIngestions($substanceFilter; $consumerFilter) as $ingestions |
  ($ingestions | ingestionsByConsumer) as $ingestionsByConsumer |

  ($ingestionsByConsumer | keys) as $consumerNames |

  "" as $experienceStatsText |
  $experienceStatsText | 
    if $withTitle then 
      . += ($experience | formatExperienceTitle | . + "\n")
    end | . as $experienceStatsText |

  reduce $consumerNames[] as $consumerName ($experienceStatsText;
    . as $experienceStatsText |
    $experienceStatsText |
      if ($consumerNames != ["default"]) 
        then . += "Consumer: \($consumerName)\n"
      end | . as $experienceStatsText  |

    ($stats.[$consumerName] | keys) as $substanceNames |


    $experienceStatsText | reduce $substanceNames[] as $substanceName (.; 
      . as $experienceStatsText |

      ($stats.[$consumerName].[$substanceName] | keys) as $ingestionMethods |

      ($experienceStatsText | . += "Substance: \($substanceName)\n") as $experienceStatsText |

      reduce ($stats.[$consumerName].[$substanceName] | to_entries)[] as $substanceStats ($experienceStatsText;
        . as $experienceStatsText |

        $substanceStats |
        .key as $ingestionMethod |
        ifNullDefault(.value.dose; "Unknown") as $dose |
        .value.unit as $unit |

        ($experienceStatsText | . += "Dose (\($ingestionMethod | titleCase)): \($dose) \($unit)\n") as $experienceStatsText | 
        $experienceStatsText
      ) | . as $experienceStatsText |

      $experienceStatsText | if ($ingestionMethods | length > 1) then 
        ($stats | calculateCombinedDose($substanceName; $consumerName)) as $combinedDose |
        . += "Combined Dose: \(ifNullDefault($combinedDose.dose; "Unknown")) \($combinedDose.unit)\n" 
      end | . as $experienceStatsText |

      $experienceStatsText | . += "\n"
    )
  ) | rtrimstr("\n\n");

def printExperienceLog($customUnits; $substanceFilter; $consumerFilter; $pretty; $withTitle):
  . as $experience |
  ($consumerFilter // ["default"]) as $consumerFilter |

  $experience.ingestions | sort_by(.sortDate) | filterIngestions($substanceFilter; $consumerFilter) as $ingestions |

  $ingestions | ingestionsConsumerNames as $consumerNames |

  if ($consumerNames == ["default"])
  then
    ["Substance", "Dose", "ROA", "Consumer", "Time"]
  else 
    ["Substance", "Dose", "ROA", "Time"]
  end | . as $columnTitles |

  (reduce $ingestions[] as $ingestion ([]; . as $rows |
    $ingestion |
    .substanceName as $substanceName |
    ifNullDefault(.consumerName; "default") as $consumerName |
    formatIngestionDose($customUnits) as $doseText |
    formatIngestionROA($customUnits) as $roaText |
    formatIngestionTime as $timeText |
    $rows | . += [
      if ($consumerNames != ["default"])
      then [$substanceName, $doseText, $roaText, $consumerName, $timeText]
      else [$substanceName, $doseText, $roaText, $timeText] end
    ]
  )) as $rows |

  if $pretty then
    printPrettyTable(
      if $withTitle then ($experience | formatExperienceTitle) else null end;
      $columnTitles;
      $rows
    )
  else
    $rows | map(join(" | ")) | rtrimstr(" ") as $rows |
    (if $withTitle then [$experience | formatExperienceTitle] else [] end) + $rows | join("\n")
  end | . as $ingestionLog |
  
  $ingestionLog;

def printExperiencesAdvanced($customUnits; $substanceFilter; $consumerFilter; $sortMethod; $sortOptions; $sortFilterString; sortFilter):
  . as $experiences |
  ($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 sortFilterExperiences:
    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 |
      (reduce $experiencesData[] as $experienceData ([]; 
        ($experienceData.experience.ingestions) as $ingestions |
        . += 
            if ($experienceData.experience.ingestions |
              any(
                .substanceName as $substanceName |
                $substanceFilter |
                  any(index($substanceName))
              )
            ) then [$experienceData] else [] end
      ));

    def filterByConsumerFilter:
      . as $experiencesData |
      (reduce $experiencesData[] as $experienceData ([]; 
        ($experienceData.experience.ingestions) as $ingestions |
        . += 
            if ($experienceData.experience.ingestions |
              any(
                ifNullDefault(.consumerName; "default") 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
                    ifNullDefault(.consumerName; "default") == $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
      assert($sortOptions.substanceName != null; "substanceName not provided as sortOption") |
      filterBySubstanceAndConsumer | 
        sort_by(highestCombinedDoseSort, oldToNewSort) |
          reverse
    elif 
      $sortMethod == "highest-dose-for-method"
    then
      assert($sortOptions.substanceName != null; "substanceName not provided as sortOption") |
      filterByIngestionMethodForSubstance | 
        sort_by(highestMethodDoseSort, oldToNewSort) |
          reverse
    end;

  def sortFilterFromString($filterString):
    (reduce ($filterString | split("|"))[] as $filter (.; 
      if 
        $filter == "reverse" then . | reverse
      elif ($filter | startswith("firstN")) then
        ($filter | ltrimstr("firstN(") | rtrimstr(")")) as $arg |
          . | firstN($arg | try fromjson catch error("invalid number passed to firstN"))
      end
    ));

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

  $experiences |
    experiencesWithExtraData($customUnits) |
    sortFilterExperiences |
    if ($sortFilterString != null) then . | sortFilterFromString($sortFilterString) end | 
    if (sortFilter != null) then . | sortFilter end |
    .[] as $entry |
      $entry.experience as $experience |
      $entry.stats as $stats |
      ($experience | printExperienceLog(
        $customUnits; 
        $substanceFilter; 
        $consumerFilter; 
        true;
        true
      )) + 
      "\nCumulative Doses:\n" +
      ($experience | printExperienceStats(
        $stats; 
        $substanceFilter;
        $consumerFilter;
        false
      )) + "\n";


def main($ARGS):
  def usage:
    [
      "psychonaut_journal_stats {printExperience,printExperiencesAdvanced}",
      ""
    ] | join("\n") | halt_error(1);

  . as $exportData |
  $exportData | types::ensureExportData |

  ($ARGS | parseArgs) as $parsedArgs |

  $parsedArgs.nonArgs[0] as $program |
  ($parsedArgs | .nonArgs |= $parsedArgs.nonArgs[1:]) as $parsedArgs |

  if $program == null then
    if any($parsedArgs.shortArgs[]; . == "h") then usage end |
    if ($parsedArgs.longArgs | has("help")) then usage end |

    usage
  elif $program == "printExperience" then 
    def printExperienceUsage($reason):
      [
        $reason,
        "Usage: printExperience [experienceTitle] --pretty=bool --title=bool --stats=bool --substance-filter=[substanceNames,] --consumer-filter=[consumerNames,]",
        ""
      ] | map(select(. != null))  | join("\n") | halt_error(1);

    if any($parsedArgs.shortArgs[]; . == "h") then printExperienceUsage(null) end |
    if ($parsedArgs.longArgs | has("help")) then printExperienceUsage(null) end |

    $parsedArgs.nonArgs[0] as $experienceTitle |
    if $experienceTitle == null then printExperienceUsage("experienceTitle not provided") end |

    {
      substanceFilter: null,
      consumerFilter: null,
      pretty: true,
      withTitle: true,
      withStats: true,
    } as $defaultOptions |
    $defaultOptions as $options |

    reduce ($parsedArgs.longArgs | to_entries[]) as $longArg ($options; (
      $longArg.key as $arg |
      $longArg.value as $value |

      if $arg == "pretty" then
        .pretty |= (ifNullDefault($value; $defaultOptions.pretty) | parseArgBool)
      end |
      if $arg == "title" then
        .withTitle |= (ifNullDefault($value; $defaultOptions.withTitle) | parseArgBool)
      end |
      if $arg == "stats" then
        .withStats |= (ifNullDefault($value; $defaultOptions.withStats) | parseArgBool)
      end |
      if $arg == "substance-filter" then
        .substanceFilter |= ($parsedArgs.longArgs.["substance-filter"] | split(","))
      end |
      if $arg == "consumer-filter" then
        .consumerFilter |= ($parsedArgs.longArgs.["consumer-filter"] | split(","))
      end
    )) | . as $options |

    $exportData.experiences | experienceByTitle($experienceTitle) as $experience |
    if $experience == null then error("Experience not found") end | 

    ($experience | printExperienceLog(
      $exportData.customUnits;
      $options.substanceFilter;
      $options.consumerFilter;
      $options.pretty;
      $options.withTitle
    )) as $ingestionLog | 
    
    if $options.withStats then 
      ($experience | experienceStats($exportData.customUnits)) as $stats |
      $ingestionLog + 
        "\n\nCumulative Doses:\n" +
        ($experience | printExperienceStats(
          $stats; 
          $options.substanceFilter;
          $options.consumerFilter;
          false
        ))
    else 
      $ingestionLog
    end
  elif $program == "printExperiencesAdvanced" then 
    def printExperiencesAdvancedUsage($reason):
      [
        $reason,
        "Usage: printExperiencesAdvanced --substance-filter=[substanceNames,]  --consumer-filter=[consumerNames,] --sort-method={old-to-new,highest-combined-dose,highest-dose-for-method} --sort-option-substance-name=[string] --sort-options-ingestion-method=[string] --sort-options-consumer-name=[string]  --sort-filter={firstN(x),reverse}",
        ""
      ] | map(select(. != null))  | join("\n") | halt_error(1);

    if any($parsedArgs.shortArgs[]; . == "h") then printExperiencesAdvancedUsage(null) end |
    if ($parsedArgs.longArgs | has("help")) then printExperiencesAdvancedUsage(null) end |

    {
      substanceFilter: null,
      consumerFilter: null,
      sortMethod: "old-to-new",
      sortOptions: {},
      sortFilter: null
    } as $defaultOptions |
    $defaultOptions as $options |

    reduce ($parsedArgs.longArgs | to_entries[]) as $longArg ($options; (
      $longArg.key as $arg |
      $longArg.value as $value |
      if $arg == "substance-filter" then
        .substanceFilter |= ($parsedArgs.longArgs.["substance-filter"] | split(","))
      end |
      if $arg == "consumer-filter" then
        .consumerFilter |= ($parsedArgs.longArgs.["consumer-filter"] | split(","))
      end |
      if $arg == "sort-method" then
        .sortMethod |= $parsedArgs.longArgs.["sort-method"]
      end |
      if $arg == "sort-option-substance-name" then
        .sortOptions.substanceName |= $parsedArgs.longArgs.["sort-option-substance-name"]
      end |
      if $arg == "sort-option-ingestion-method" then
        .sortOptions.ingestionMethod |= $parsedArgs.longArgs.["sort-option-ingestion-method"]
      end |
      if $arg == "sort-option-consumer-name" then
        .sortOptions.consumerName |= $parsedArgs.longArgs.["sort-option-consumer-name"]
      end |
      if $arg == "sort-filter" then
        .sortFilter |= $parsedArgs.longArgs.["sort-filter"]
      end
    )) | . as $options |

    $exportData.experiences |
      printExperiencesAdvanced(
        $exportData.customUnits;
        $options.substanceFilter;
        $options.consumerFilter;
        $options.sortMethod; 
        $options.sortOptions;
        $options.sortFilter; null
      )
  else
    usage
  end;

main($ARGS)