From 58f4a16569e75331c38ee7dbc0595ba5ebb76db9 Mon Sep 17 00:00:00 2001
From: Nick Moore <nick@zoic.org>
Date: Fri, 28 Jun 2024 18:48:30 +1000
Subject: [PATCH] start on curve fit plugin

---
 countess/core/parameters.py   | 35 ++++++++++++++++++++++++++++++++++-
 countess/plugins/curve_fit.py | 23 +++++++++++++++++++++++
 pyproject.toml                |  2 ++
 3 files changed, 59 insertions(+), 1 deletion(-)
 create mode 100644 countess/plugins/curve_fit.py

diff --git a/countess/core/parameters.py b/countess/core/parameters.py
index bdf9d9d..0da0fad 100644
--- a/countess/core/parameters.py
+++ b/countess/core/parameters.py
@@ -2,13 +2,23 @@
 import math
 import os.path
 import re
-from typing import Any, Iterable, Mapping, Optional, Type, Union
+from typing import Any, Dict, Iterable, List, Mapping, Optional, Type, Union
 
 import pandas as pd
 
+from countess.utils.pandas import get_all_columns
+
 PARAM_DIGEST_HASH = "sha256"
 
 
+def make_prefix_groups(strings: List[str]) -> Dict[str, List[str]]:
+    groups : Dict[str, List[str]] = {}
+    for s in strings:
+        if m := re.match(r'(.*?)_+([^_]+)$', s):
+            groups.setdefault(m.group(1), []).append(s)
+    return { k: v for k, v in groups.items() if len(v) > 1 }
+
+
 class BaseParam:
     """Represents the parameters which can be set on a plugin."""
 
@@ -387,6 +397,29 @@ def get_column(self, df):
             return _dataframe_get_column(df, self.value)
 
 
+class ColumnGroupChoiceParam(ChoiceParam):
+
+    def set_column_choices(self, choices):
+        self.set_choices(make_prefix_groups(choices))
+
+    def get_column_names(self, df):
+        column_names = get_all_columns(df).keys()
+        return [ n for n in column_names if n.startswith(self.value) ]
+
+
+class ColumnGroupOrNoneChoiceParam(ColumnGroupChoiceParam):
+    DEFAULT_VALUE = "— NONE —"
+
+    def set_choices(self, choices: Iterable[str]):
+        super().set_choices([self.DEFAULT_VALUE] + list(choices))
+
+    def is_none(self):
+        return self.value == self.DEFAULT_VALUE
+
+    def is_not_none(self):
+        return self.value != self.DEFAULT_VALUE
+
+
 class ColumnOrIndexChoiceParam(ColumnChoiceParam):
     DEFAULT_VALUE = "— INDEX —"
 
diff --git a/countess/plugins/curve_fit.py b/countess/plugins/curve_fit.py
new file mode 100644
index 0000000..a209b71
--- /dev/null
+++ b/countess/plugins/curve_fit.py
@@ -0,0 +1,23 @@
+from scipy.optimize import curve_fit
+
+from countess import VERSION
+from countess.core.plugins import PandasTransformDictToDictPlugin
+from countess.core.parameters import ChoiceParam, ColumnGroupChoiceParam, ColumnGroupOrNoneChoiceParam
+from countess.core.logger import Logger
+
+FUNCTIONS = [ 'foo', 'bar', 'baz' ]
+
+class CurveFitPlugin(PandasTransformDictToDictPlugin):
+    name = "Curve Fit Plugin"
+    description = "Fit rows of data to curves"
+
+    version = VERSION
+
+    parameters = {
+        "xaxis": ColumnGroupOrNoneChoiceParam("X Axis", None, []),
+        "yaxis": ColumnGroupChoiceParam("Y Axis", None, []),
+        "function": ChoiceParam("Function", FUNCTIONS[0], FUNCTIONS),
+    }
+
+    def process_dict(self, data: dict, logger: Logger):
+        return data
diff --git a/pyproject.toml b/pyproject.toml
index 8d1a90c..45ca87d 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -26,6 +26,7 @@ dependencies = [
     'pandas~=2.1.0',
     'psutil~=5.9.5',
     'rapidfuzz~=2.15.1',
+    'scipy~=1.14.0',
     'tkinterweb~=3.23.5',
     'ttkthemes~=3.2.2',
     'typing_extensions~=4.8.0',
@@ -64,6 +65,7 @@ data_table = "countess.plugins.data_table:DataTablePlugin"
 correlation = "countess.plugins.correlation:CorrelationPlugin"
 hgvs_parser = "countess.plugins.hgvs_parser:HgvsParserPlugin"
 column = "countess.plugins.column:ColumnToolPlugin"
+curve_fit = "countess.plugins.curve_fit:CurveFitPlugin"
 
 [project.entry-points.gui_scripts]
 countess_gui = "countess.gui.main:main"