From dac6919fcd089ce868ad5b9ed1065843a4fb7966 Mon Sep 17 00:00:00 2001
From: Starbeamrainbowlabs <sbrl@starbeamrainbowlabs.com>
Date: Mon, 25 Jul 2022 19:00:53 +0100
Subject: [PATCH] ai: start creating initial scaffolding

---
 aimodel/src/index.py      | 33 +++++++++++++++++++++++++++++++
 aimodel/src/parse_args.py | 41 +++++++++++++++++++++++++++++++++++++++
 2 files changed, 74 insertions(+)
 create mode 100644 aimodel/src/parse_args.py

diff --git a/aimodel/src/index.py b/aimodel/src/index.py
index e69de29..7eb9007 100644
--- a/aimodel/src/index.py
+++ b/aimodel/src/index.py
@@ -0,0 +1,33 @@
+#!/usr/bin/env python3
+
+import importlib
+import sys
+
+from loguru import logger
+
+def init_logging():
+	pass
+
+from parse_args import parse_args
+
+def main():
+	subcommand, args = parse_args()
+	if args == None:
+		return
+	
+	imported_module = importlib.import_module(f"subcommands.{subcommand}")
+	# TODO: Support multiple subcommands here
+	match subcommand:
+		case "pretrain":
+			imported_module.pretrain(args)
+		case _:
+			sys.stderr.write(f"Error: The subcommand '{subcommand}' hasn't been registered in index.py yet.\nThis is a bug.\n")
+			exit(1)
+	
+	
+
+if __name__ == "__main__":
+    main()
+else:
+    print("This script must be run directly. It cannot be imported.")
+    exit(1)
diff --git a/aimodel/src/parse_args.py b/aimodel/src/parse_args.py
new file mode 100644
index 0000000..957ca86
--- /dev/null
+++ b/aimodel/src/parse_args.py
@@ -0,0 +1,41 @@
+
+import argparse
+import sys
+import re
+import importlib
+
+# import pysnooper
+
+# @pysnooper.snoop()
+def parse_args():
+	"""Defines and parses the CLI arguments."""
+	
+	
+	if len(sys.argv) < 2:
+		sys.stderr.write("""
+This program trains, manipulates, visualises, and manages a contrastive learning based rainfall radar → water depth prediction model.
+It functions by first finding relationships between the rainfall radar data and the water depth + heightmap data (the 'pretrain' subcommand). After this, a decoder model to predict water depth (modelled as an image segmentation task), can then be trained.
+
+Available subcommands:
+	pretrain	Pretrain a contrastive learning model as an encoder.
+
+For more information, do src/index.py <subcommand> --help.
+""")
+		exit(0)
+	
+	subcommand = re.sub(r'[^a-z0-9-]', '', sys.argv[1])
+	
+	subcommand_argparser = importlib.import_module(f"subcommands.{subcommand}").parse_args
+	
+	parser = subcommand_argparser()
+	# sys.stderr.write(f"Error: Unknown subcommand '{subcommand} (try --help).\n")
+	# exit(1)
+	if parser == None:
+		sys.stderr.write(f"Error: The subcommand '{subcommand}' did not return an argument parser. This is a bug.\n")
+		exit(1)
+	
+	parser.add_argument("--only-gpu",
+		help="If the GPU is not available, exit with an error (useful on shared HPC systems to avoid running out of memory & affecting other users)", action="store_true")
+	
+	return subcommand, parser.parse_args(args=sys.argv[2:])
+