From 7d1e6dc6067249b7c252ee37e13d9f826598e9ec Mon Sep 17 00:00:00 2001
From: lwark <lando.wark@gmail.com>
Date: Thu, 21 Nov 2024 08:46:41 -0600
Subject: [PATCH] Mid-progress adding controls to pydantic creation.

---
 CHANGELOG.md                                  |   4 +
 TODO.md                                       |   2 +-
 src/submissions/backend/db/__init__.py        |  42 +++++-
 src/submissions/backend/db/models/__init__.py |  22 ++-
 src/submissions/backend/db/models/controls.py |   9 ++
 src/submissions/backend/db/models/kits.py     |   7 +-
 .../backend/db/models/submissions.py          |  31 +++--
 .../backend/validators/__init__.py            |   2 +-
 src/submissions/backend/validators/pydant.py  |  35 ++++-
 src/submissions/frontend/widgets/app.py       |  11 +-
 src/submissions/frontend/widgets/misc.py      |  14 +-
 .../frontend/widgets/omni_search.py           |  60 +++++---
 .../frontend/widgets/sample_search.py         | 128 ------------------
 .../frontend/widgets/submission_details.py    |  15 +-
 .../frontend/widgets/submission_widget.py     |  21 +--
 src/submissions/tools/__init__.py             |   8 ++
 16 files changed, 224 insertions(+), 187 deletions(-)
 delete mode 100644 src/submissions/frontend/widgets/sample_search.py

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 1977cf1..fa30cf6 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -1,3 +1,7 @@
+## 202411.04
+
+- Add reagent from scrape now limits roles to those found in kit to prevent confusion.
+
 ## 202411.01
 
 - Code clean up.
diff --git a/TODO.md b/TODO.md
index 6861ed5..d743833 100644
--- a/TODO.md
+++ b/TODO.md
@@ -1,4 +1,4 @@
-- [ ] Find a way to merge omni_search and sample_search
+- [x] Find a way to merge omni_search and sample_search
 - [x] Allow parsing of custom fields to a json 'custom' field in _basicsubmissions
 - [x] Upgrade to generators when returning lists.
 - [x] Revamp frontend.widgets.controls_chart to include visualizations?
diff --git a/src/submissions/backend/db/__init__.py b/src/submissions/backend/db/__init__.py
index 1bd444d..f73b413 100644
--- a/src/submissions/backend/db/__init__.py
+++ b/src/submissions/backend/db/__init__.py
@@ -1,8 +1,10 @@
 """
 All database related operations.
 """
-from sqlalchemy import event
+import sqlalchemy.orm
+from sqlalchemy import event, inspect
 from sqlalchemy.engine import Engine
+
 from tools import ctx
 
 
@@ -16,7 +18,7 @@ def set_sqlite_pragma(dbapi_connection, connection_record):
     Args:
         dbapi_connection (_type_): _description_
         connection_record (_type_): _description_
-    """    
+    """
     cursor = dbapi_connection.cursor()
     # print(ctx.database_schema)
     if ctx.database_schema == "sqlite":
@@ -34,3 +36,39 @@ def set_sqlite_pragma(dbapi_connection, connection_record):
 
 
 from .models import *
+
+
+def update_log(mapper, connection, target):
+    logger.debug("\n\nBefore update\n\n")
+    state = inspect(target)
+    logger.debug(state)
+    update = dict(user=getuser(), time=datetime.now(), object=str(state.object), changes=[])
+    logger.debug(update)
+    for attr in state.attrs:
+        hist = attr.load_history()
+        if not hist.has_changes():
+            continue
+        added = [str(item) for item in hist.added]
+        deleted = [str(item) for item in hist.deleted]
+        change = dict(field=attr.key, added=added, deleted=deleted)
+        logger.debug(f"Adding: {pformat(change)}")
+        try:
+            update['changes'].append(change)
+        except Exception as e:
+            logger.error(f"Something went horribly wrong adding attr: {attr.key}: {e}")
+            continue
+
+    logger.debug(f"Adding to audit logs: {pformat(update)}")
+    if update['changes']:
+        # Note: must use execute as the session will be busy at this point.
+        # https://medium.com/@singh.surbhicse/creating-audit-table-to-log-insert-update-and-delete-changes-in-flask-sqlalchemy-f2ca53f7b02f
+        table = AuditLog.__table__
+        logger.debug(f"Adding to {table}")
+        connection.execute(table.insert().values(**update))
+        # logger.debug("Here is where I would insert values, if I was able.")
+    else:
+        logger.info(f"No changes detected, not updating logs.")
+
+
+# event.listen(LogMixin, 'after_update', update_log, propagate=True)
+# event.listen(LogMixin, 'after_insert', update_log, propagate=True)
diff --git a/src/submissions/backend/db/models/__init__.py b/src/submissions/backend/db/models/__init__.py
index e7cd746..5e4c218 100644
--- a/src/submissions/backend/db/models/__init__.py
+++ b/src/submissions/backend/db/models/__init__.py
@@ -5,7 +5,7 @@ from __future__ import annotations
 import sys, logging
 
 import pandas as pd
-from sqlalchemy import Column, INTEGER, String, JSON
+from sqlalchemy import Column, INTEGER, String, JSON, event, inspect
 from sqlalchemy.orm import DeclarativeMeta, declarative_base, Query, Session
 from sqlalchemy.ext.declarative import declared_attr
 from sqlalchemy.exc import ArgumentError
@@ -22,6 +22,12 @@ Base: DeclarativeMeta = declarative_base()
 logger = logging.getLogger(f"submissions.{__name__}")
 
 
+class LogMixin(Base):
+
+    __abstract__ = True
+
+
+
 class BaseClass(Base):
     """
     Abstract class to pass ctx values to all SQLAlchemy objects.
@@ -99,6 +105,15 @@ class BaseClass(Base):
         singles = list(set(cls.singles + BaseClass.singles))
         return dict(singles=singles)
 
+    @classmethod
+    def find_regular_subclass(cls, name: str | None = None):
+        if not name:
+            return cls
+        if " " in name:
+            search = name.title().replace(" ", "")
+        logger.debug(f"Searching for subclass: {search}")
+        return next((item for item in cls.__subclasses__() if item.__name__ == search), cls)
+
     @classmethod
     def fuzzy_search(cls, **kwargs):
         query: Query = cls.__database_session__.query(cls)
@@ -175,9 +190,11 @@ class BaseClass(Base):
         """
         # logger.debug(f"Saving object: {pformat(self.__dict__)}")
         report = Report()
+        state = inspect(self)
         try:
             self.__database_session__.add(self)
             self.__database_session__.commit()
+            return state
         except Exception as e:
             logger.critical(f"Problem saving object: {e}")
             logger.error(f"Error message: {type(e)}")
@@ -186,6 +203,8 @@ class BaseClass(Base):
             return report
 
 
+
+
 class ConfigItem(BaseClass):
     """
     Key:JSON objects to store config settings in database. 
@@ -222,6 +241,7 @@ from .controls import *
 from .organizations import *
 from .kits import *
 from .submissions import *
+from .audit import AuditLog
 
 # NOTE: Add a creator to the submission for reagent association. Assigned here due to circular import constraints.
 # https://docs.sqlalchemy.org/en/20/orm/extensions/associationproxy.html#sqlalchemy.ext.associationproxy.association_proxy.params.creator
diff --git a/src/submissions/backend/db/models/controls.py b/src/submissions/backend/db/models/controls.py
index cb31e85..bf62552 100644
--- a/src/submissions/backend/db/models/controls.py
+++ b/src/submissions/backend/db/models/controls.py
@@ -16,6 +16,7 @@ from typing import List, Literal, Tuple, Generator
 from dateutil.parser import parse
 from re import Pattern
 
+
 logger = logging.getLogger(f"submissions.{__name__}")
 
 
@@ -429,6 +430,10 @@ class PCRControl(Control):
         fig = PCRFigure(df=df, modes=[])
         return report, fig
 
+    def to_pydantic(self):
+        from backend.validators import PydPCRControl
+        return PydPCRControl(**self.to_sub_dict())
+
 
 class IridaControl(Control):
 
@@ -878,3 +883,7 @@ class IridaControl(Control):
             exclude = [re.sub(rerun_regex, "", sample) for sample in sample_names if rerun_regex.search(sample)]
             df = df[df.name not in exclude]
         return df
+
+    def to_pydantic(self):
+        from backend.validators import PydIridaControl
+        return PydIridaControl(**self.__dict__)
diff --git a/src/submissions/backend/db/models/kits.py b/src/submissions/backend/db/models/kits.py
index fda0a5d..c57f7b2 100644
--- a/src/submissions/backend/db/models/kits.py
+++ b/src/submissions/backend/db/models/kits.py
@@ -567,7 +567,7 @@ class Reagent(BaseClass):
         return cls.execute_query(query=query, limit=limit)
 
     @check_authorization
-    def edit_from_search(self, **kwargs):
+    def edit_from_search(self, obj, **kwargs):
         from frontend.widgets.misc import AddReagentForm
         role = ReagentRole.query(kwargs['role'])
         if role:
@@ -1279,7 +1279,10 @@ class SubmissionReagentAssociation(BaseClass):
         Returns:
             str: Representation of this SubmissionReagentAssociation
         """
-        return f"<{self.submission.rsl_plate_num} & {self.reagent.lot}>"
+        try:
+            return f"<{self.submission.rsl_plate_num} & {self.reagent.lot}>"
+        except AttributeError:
+            return f"<Unknown Submission & {self.reagent.lot}"
 
     def __init__(self, reagent=None, submission=None):
         if isinstance(reagent, list):
diff --git a/src/submissions/backend/db/models/submissions.py b/src/submissions/backend/db/models/submissions.py
index 3dce4fe..98b2067 100644
--- a/src/submissions/backend/db/models/submissions.py
+++ b/src/submissions/backend/db/models/submissions.py
@@ -12,8 +12,8 @@ from zipfile import ZipFile
 from tempfile import TemporaryDirectory, TemporaryFile
 from operator import itemgetter
 from pprint import pformat
-from . import BaseClass, Reagent, SubmissionType, KitType, Organization, Contact
-from sqlalchemy import Column, String, TIMESTAMP, INTEGER, ForeignKey, JSON, FLOAT, case
+from . import BaseClass, Reagent, SubmissionType, KitType, Organization, Contact, LogMixin
+from sqlalchemy import Column, String, TIMESTAMP, INTEGER, ForeignKey, JSON, FLOAT, case, event, inspect
 from sqlalchemy.orm import relationship, validates, Query
 from sqlalchemy.orm.attributes import flag_modified
 from sqlalchemy.ext.associationproxy import association_proxy
@@ -36,7 +36,7 @@ from PIL import Image
 logger = logging.getLogger(f"submissions.{__name__}")
 
 
-class BasicSubmission(BaseClass):
+class BasicSubmission(BaseClass, LogMixin):
     """
     Concrete of basic submission which polymorphs into BacterialCulture and Wastewater
     """
@@ -544,7 +544,7 @@ class BasicSubmission(BaseClass):
                     field_value = len(self.samples)
                 else:
                     field_value = value
-            case "ctx" | "csv" | "filepath" | "equipment":
+            case "ctx" | "csv" | "filepath" | "equipment" | "controls":
                 return
             case item if item in self.jsons():
                 match key:
@@ -577,10 +577,13 @@ class BasicSubmission(BaseClass):
                 except AttributeError:
                     field_value = value
         # NOTE: insert into field
-        try:
-            self.__setattr__(key, field_value)
-        except AttributeError as e:
-            logger.error(f"Could not set {self} attribute {key} to {value} due to \n{e}")
+        current = self.__getattribute__(key)
+        if field_value and current != field_value:
+            logger.debug(f"Updated value: {key}: {current} to {field_value}")
+            try:
+                self.__setattr__(key, field_value)
+            except AttributeError as e:
+                logger.error(f"Could not set {self} attribute {key} to {value} due to \n{e}")
 
     def update_subsampassoc(self, sample: BasicSample, input_dict: dict):
         """
@@ -1339,7 +1342,7 @@ class BasicSubmission(BaseClass):
 
 # Below are the custom submission types
 
-class BacterialCulture(BasicSubmission):
+class BacterialCulture(BasicSubmission, LogMixin):
     """
     derivative submission type from BasicSubmission
     """
@@ -1426,7 +1429,7 @@ class BacterialCulture(BasicSubmission):
         return input_dict
 
 
-class Wastewater(BasicSubmission):
+class Wastewater(BasicSubmission, LogMixin):
     """
     derivative submission type from BasicSubmission
     """
@@ -2189,6 +2192,8 @@ class BasicSample(BaseClass):
     Base of basic sample which polymorphs into BCSample and WWSample
     """
 
+    searchables = ['submitter_id']
+
     id = Column(INTEGER, primary_key=True)  #: primary key
     submitter_id = Column(String(64), nullable=False, unique=True)  #: identification from submitter
     sample_type = Column(String(32))  #: mode_sub_type of sample
@@ -2509,6 +2514,9 @@ class BasicSample(BaseClass):
         if dlg.exec():
             pass
 
+    def edit_from_search(self, obj, **kwargs):
+        self.show_details(obj)
+
 
 # Below are the custom sample types
 
@@ -2516,6 +2524,9 @@ class WastewaterSample(BasicSample):
     """
     Derivative wastewater sample
     """
+
+    searchables = BasicSample.searchables + ['ww_processing_num', 'ww_full_sample_id', 'rsl_number']
+
     id = Column(INTEGER, ForeignKey('_basicsample.id'), primary_key=True)
     ww_processing_num = Column(String(64))  #: wastewater processing number
     ww_full_sample_id = Column(String(64))  #: full id given by entrics
diff --git a/src/submissions/backend/validators/__init__.py b/src/submissions/backend/validators/__init__.py
index c390287..19686b1 100644
--- a/src/submissions/backend/validators/__init__.py
+++ b/src/submissions/backend/validators/__init__.py
@@ -203,4 +203,4 @@ class RSLNamer(object):
 
 
 from .pydant import PydSubmission, PydKit, PydContact, PydOrganization, PydSample, PydReagent, PydReagentRole, \
-    PydEquipment, PydEquipmentRole, PydTips
+    PydEquipment, PydEquipmentRole, PydTips, PydPCRControl, PydIridaControl
diff --git a/src/submissions/backend/validators/pydant.py b/src/submissions/backend/validators/pydant.py
index d394871..0a0a44f 100644
--- a/src/submissions/backend/validators/pydant.py
+++ b/src/submissions/backend/validators/pydant.py
@@ -786,9 +786,10 @@ class PydSubmission(BaseModel, extra='allow'):
         """
         report = Report()
         dicto = self.improved_dict()
+        logger.warning(f"\n\nQuery or create: {self.submission_type['value']}, {self.rsl_plate_num['value']}")
         instance, result = BasicSubmission.query_or_create(submission_type=self.submission_type['value'],
                                                            rsl_plate_num=self.rsl_plate_num['value'])
-        logger.debug(f"Result of query or create: {type(result)}")
+        logger.debug(f"Result of query or create: {instance}")
         report.add_result(result)
         self.handle_duplicate_samples()
         # logger.debug(f"Here's our list of duplicate removed samples: {self.samples}")
@@ -830,7 +831,7 @@ class PydSubmission(BaseModel, extra='allow'):
                         equip, association = equip.toSQL(submission=instance)
                         if association is not None:
                             instance.submission_equipment_associations.append(association)
-                    logger.debug(f"Equipment associations:\n\n{pformat(instance.submission_equipment_associations)}")
+                    logger.debug(f"Equipment associations: {instance.submission_equipment_associations}")
                 case "tips":
                     for tips in self.tips:
                         if tips is None:
@@ -871,16 +872,18 @@ class PydSubmission(BaseModel, extra='allow'):
                 case _:
                     try:
                         instance.set_attribute(key=key, value=value)
+                        # instance.update({key:value})
                     except AttributeError as e:
                         logger.error(f"Could not set attribute: {key} to {value} due to: \n\n {e}")
                         continue
                     except KeyError:
                         continue
+        print(f"\n\n{instance}\n\n")
         try:
             # logger.debug(f"Calculating costs for procedure...")
             instance.calculate_base_cost()
         except (TypeError, AttributeError) as e:
-            # logger.debug(f"Looks like that kit doesn't have cost breakdown yet due to: {e}, using full plate cost.")
+            logger.debug(f"Looks like that kit doesn't have cost breakdown yet due to: {e}, using 0.")
             try:
                 instance.run_cost = instance.extraction_kit.cost_per_run
             except AttributeError:
@@ -1104,3 +1107,29 @@ class PydEquipmentRole(BaseModel):
         """
         from frontend.widgets.equipment_usage import RoleComboBox
         return RoleComboBox(parent=parent, role=self, used=used)
+
+
+class PydPCRControl(BaseModel):
+    name: str
+    subtype: str
+    target: str
+    ct: float
+    reagent_lot: str
+    submitted_date: datetime  #: Date submitted to Robotics
+    submission_id: int
+    controltype_name: str
+
+
+class PydIridaControl(BaseModel, extra='ignore'):
+    name: str
+    contains: list | dict  #: unstructured hashes in contains.tsv for each organism
+    matches: list | dict  #: unstructured hashes in matches.tsv for each organism
+    kraken: list | dict  #: unstructured output from kraken_report
+    subtype: str  #: EN-NOS, MCS-NOS, etc
+    refseq_version: str  #: version of refseq used in fastq parsing
+    kraken2_version: str
+    kraken2_db_version: str
+    sample_id: int
+    submitted_date: datetime  #: Date submitted to Robotics
+    submission_id: int
+    controltype_name: str
diff --git a/src/submissions/frontend/widgets/app.py b/src/submissions/frontend/widgets/app.py
index 3149818..60c856b 100644
--- a/src/submissions/frontend/widgets/app.py
+++ b/src/submissions/frontend/widgets/app.py
@@ -13,7 +13,7 @@ from PyQt6.QtGui import QAction
 from pathlib import Path
 from markdown import markdown
 from __init__ import project_path
-from backend import SubmissionType, Reagent
+from backend import SubmissionType, Reagent, BasicSample
 from tools import check_if_app, Settings, Report, jinja_template_loading, check_authorization, page_size
 from .functions import select_save_file, select_open_file
 from datetime import date
@@ -23,7 +23,7 @@ import logging, webbrowser, sys, shutil
 from .submission_table import SubmissionsSheet
 from .submission_widget import SubmissionFormContainer
 from .controls_chart import ControlsViewer
-from .sample_search import SampleSearchBox
+# from .sample_search import SampleSearchBox
 from .summary import Summary
 from .omni_search import SearchBox
 
@@ -138,6 +138,7 @@ class App(QMainWindow):
         self.yamlImportAction.triggered.connect(self.import_ST_yaml)
         self.table_widget.pager.current_page.textChanged.connect(self.update_data)
         self.editReagentAction.triggered.connect(self.edit_reagent)
+        self.destroyed.connect(self.final_commit)
 
     def showAbout(self):
         """
@@ -186,7 +187,8 @@ class App(QMainWindow):
         """
         Create a search for samples.
         """
-        dlg = SampleSearchBox(self)
+        # dlg = SampleSearchBox(self)
+        dlg = SearchBox(self, object_type=BasicSample, extras=[])
         dlg.exec()
 
     def backup_database(self):
@@ -251,6 +253,9 @@ class App(QMainWindow):
     def update_data(self):
         self.table_widget.sub_wid.setData(page=self.table_widget.pager.page_anchor, page_size=page_size)
 
+    def final_commit(self):
+        logger.debug("Running final commit")
+        self.ctx.database_session.commit()
 
 class AddSubForm(QWidget):
 
diff --git a/src/submissions/frontend/widgets/misc.py b/src/submissions/frontend/widgets/misc.py
index 785c738..8432e8b 100644
--- a/src/submissions/frontend/widgets/misc.py
+++ b/src/submissions/frontend/widgets/misc.py
@@ -28,7 +28,7 @@ class AddReagentForm(QDialog):
     """
 
     def __init__(self, reagent_lot: str | None = None, reagent_role: str | None = None, expiry: date | None = None,
-                 reagent_name: str | None = None) -> None:
+                 reagent_name: str | None = None, kit: str | KitType | None = None) -> None:
         super().__init__()
         if reagent_name is None:
             reagent_name = reagent_role
@@ -58,8 +58,16 @@ class AddReagentForm(QDialog):
                 self.exp_input.setDate(QDate(1970, 1, 1))
         # NOTE: widget to get reagent type info
         self.type_input = QComboBox()
-        self.type_input.setObjectName('type')
-        self.type_input.addItems([item.name for item in ReagentRole.query()])
+        self.type_input.setObjectName('role')
+        if kit:
+            match kit:
+                case str():
+                    kit = KitType.query(name=kit)
+                case _:
+                    pass
+            self.type_input.addItems([item.name for item in ReagentRole.query() if kit in item.kit_types])
+        else:
+            self.type_input.addItems([item.name for item in ReagentRole.query()])
         # logger.debug(f"Trying to find index of {reagent_type}")
         # NOTE: convert input to user-friendly string?
         try:
diff --git a/src/submissions/frontend/widgets/omni_search.py b/src/submissions/frontend/widgets/omni_search.py
index 0fb4868..82da54d 100644
--- a/src/submissions/frontend/widgets/omni_search.py
+++ b/src/submissions/frontend/widgets/omni_search.py
@@ -19,17 +19,24 @@ class SearchBox(QDialog):
 
     def __init__(self, parent, object_type: Any, extras: List[str], **kwargs):
         super().__init__(parent)
-        self.object_type = object_type
-        # options = ["Any"] + [cls.__name__ for cls in self.object_type.__subclasses__()]
-        # self.sub_class = QComboBox(self)
-        # self.sub_class.setObjectName("sub_class")
-        # self.sub_class.currentTextChanged.connect(self.update_widgets)
-        # self.sub_class.addItems(options)
-        # self.sub_class.setEditable(False)
+        self.object_type = self.original_type = object_type
+        self.extras = extras
+        self.context = kwargs
+        self.layout = QGridLayout(self)
         self.setMinimumSize(600, 600)
-        # self.sub_class.setMinimumWidth(self.minimumWidth())
-        # self.layout.addWidget(self.sub_class, 0, 0)
-        self.results = SearchResults(parent=self, object_type=self.object_type, extras=extras, **kwargs)
+        options = ["Any"] + [cls.__name__ for cls in self.object_type.__subclasses__()]
+        if len(options) > 1:
+            self.sub_class = QComboBox(self)
+            self.sub_class.setObjectName("sub_class")
+            self.sub_class.addItems(options)
+            self.sub_class.currentTextChanged.connect(self.update_widgets)
+            self.sub_class.setEditable(False)
+            self.sub_class.setMinimumWidth(self.minimumWidth())
+            self.layout.addWidget(self.sub_class, 0, 0)
+        else:
+            self.sub_class = None
+        self.results = SearchResults(parent=self, object_type=self.object_type, extras=self.extras, **kwargs)
+        logger.debug(f"results: {self.results}")
         self.layout.addWidget(self.results, 5, 0)
         self.setLayout(self.layout)
         self.setWindowTitle(f"Search {self.object_type.__name__}")
@@ -40,10 +47,23 @@ class SearchBox(QDialog):
         """
         Changes form inputs based on sample type
         """
+        deletes = [item for item in self.findChildren(FieldSearch)]
+        # logger.debug(deletes)
+        for item in deletes:
+            item.setParent(None)
+        if not self.sub_class:
+            self.update_data()
+        else:
+            if self.sub_class.currentText() == "Any":
+                self.object_type = self.original_type
+            else:
+                self.object_type = self.original_type.find_regular_subclass(self.sub_class.currentText())
+        logger.debug(f"{self.object_type} searchables: {self.object_type.searchables}")
         for iii, searchable in enumerate(self.object_type.searchables):
-            self.widget = FieldSearch(parent=self, label=searchable, field_name=searchable)
-            self.layout.addWidget(self.widget, 1, 0)
-            self.widget.search_widget.textChanged.connect(self.update_data)
+            widget = FieldSearch(parent=self, label=searchable, field_name=searchable)
+            widget.setObjectName(searchable)
+            self.layout.addWidget(widget, 1+iii, 0)
+            widget.search_widget.textChanged.connect(self.update_data)
         self.update_data()
 
     def parse_form(self) -> dict:
@@ -60,11 +80,11 @@ class SearchBox(QDialog):
         """
         Shows dataframe of relevant samples.
         """
-        # logger.debug(f"Running update_data with sample type: {self.type}")
         fields = self.parse_form()
         # logger.debug(f"Got fields: {fields}")
         sample_list_creator = self.object_type.fuzzy_search(**fields)
         data = self.object_type.results_to_df(objects=sample_list_creator)
+        # Setting results moved to here from __init__ 202411118
         self.results.setData(df=data)
 
 
@@ -108,7 +128,6 @@ class SearchResults(QTableView):
         sets data in model
         """
         self.data = df
-        print(self.data)
         try:
             self.columns_of_interest = [dict(name=item, column=self.data.columns.get_loc(item)) for item in self.extras]
         except KeyError:
@@ -125,14 +144,15 @@ class SearchResults(QTableView):
 
     def parse_row(self, x):
         context = {item['name']: x.sibling(x.row(), item['column']).data() for item in self.columns_of_interest}
+        logger.debug(f"Context: {context}")
         try:
-            object = self.object_type.query(**{self.object_type.search: context[self.object_type.search]})
+            # object = self.object_type.query(**{self.object_type.searchables: context[self.object_type.searchables]})
+            object = self.object_type.query(**context)
         except KeyError:
             object = None
         try:
-            object.edit_from_search(**context)
-        except AttributeError:
-            pass
+            object.edit_from_search(obj=self.parent, **context)
+        except AttributeError as e:
+            logger.error(f"Error getting object function: {e}")
         self.doubleClicked.disconnect()
         self.parent.update_data()
-
diff --git a/src/submissions/frontend/widgets/sample_search.py b/src/submissions/frontend/widgets/sample_search.py
deleted file mode 100644
index 6300b5b..0000000
--- a/src/submissions/frontend/widgets/sample_search.py
+++ /dev/null
@@ -1,128 +0,0 @@
-'''
-Search box that performs fuzzy search for samples
-'''
-from pprint import pformat
-from typing import Tuple
-from pandas import DataFrame
-from PyQt6.QtCore import QSortFilterProxyModel
-from PyQt6.QtWidgets import (
-    QLabel, QVBoxLayout, QDialog,
-    QComboBox, QTableView, QWidget, QLineEdit, QGridLayout
-)
-from backend.db.models import BasicSample
-from .submission_table import pandasModel
-import logging
-
-logger = logging.getLogger(f"submissions.{__name__}")
-
-
-class SampleSearchBox(QDialog):
-
-    def __init__(self, parent):
-        super().__init__(parent)
-        self.layout = QGridLayout(self)
-        self.sample_type = QComboBox(self)
-        self.sample_type.setObjectName("sample_type")
-        self.sample_type.currentTextChanged.connect(self.update_widgets)
-        options = ["Any"] + [cls.__mapper_args__['polymorphic_identity'] for cls in BasicSample.__subclasses__()]
-        self.sample_type.addItems(options)
-        self.sample_type.setEditable(False)
-        self.setMinimumSize(600, 600)
-        self.sample_type.setMinimumWidth(self.minimumWidth())
-        self.layout.addWidget(self.sample_type, 0, 0)
-        self.results = SearchResults()
-        self.layout.addWidget(self.results, 5, 0)
-        self.setLayout(self.layout)
-        self.update_widgets()
-        self.update_data()
-
-    def update_widgets(self):
-        """
-        Changes form inputs based on sample type
-        """        
-        deletes = [item for item in self.findChildren(FieldSearch)]
-        # logger.debug(deletes)
-        for item in deletes:
-            item.setParent(None)
-        if self.sample_type.currentText() == "Any":
-            self.type = BasicSample
-        else:
-            self.type = BasicSample.find_polymorphic_subclass(self.sample_type.currentText())
-        # logger.debug(f"Sample type: {self.type}")
-        searchables = self.type.get_searchables()
-        start_row = 1
-        for iii, item in enumerate(searchables):
-            widget = FieldSearch(parent=self, label=item['label'], field_name=item['field'])
-            self.layout.addWidget(widget, start_row+iii, 0)
-            widget.search_widget.textChanged.connect(self.update_data)
-        self.update_data()
-
-    def parse_form(self) -> dict:
-        """
-        Converts form into dictionary.
-
-        Returns:
-            dict: Fields dictionary
-        """        
-        fields = [item.parse_form() for item in self.findChildren(FieldSearch)]
-        return {item[0]:item[1] for item in fields if item[1] is not None}
-
-    def update_data(self):
-        """
-        Shows dataframe of relevant samples.
-        """
-        # logger.debug(f"Running update_data with sample type: {self.type}")
-        fields = self.parse_form()
-        # logger.debug(f"Got fields: {fields}")
-        sample_list_creator = self.type.fuzzy_search(**fields)
-        data = self.type.samples_to_df(sample_list=sample_list_creator)
-        # logger.debug(f"Data: {data}")
-        self.results.setData(df=data)
-
-
-class FieldSearch(QWidget):
-
-    def __init__(self, parent, label, field_name):
-        super().__init__(parent)
-        self.layout = QVBoxLayout(self)
-        label_widget = QLabel(label)
-        self.layout.addWidget(label_widget)
-        self.search_widget = QLineEdit()
-        self.search_widget.setObjectName(field_name)
-        self.layout.addWidget(self.search_widget)
-        self.setLayout(self.layout)
-        self.search_widget.returnPressed.connect(self.enter_pressed)
-
-    def enter_pressed(self):
-        """
-        Triggered when enter is pressed on this input field.
-        """        
-        self.parent().update_data()
-
-    def parse_form(self) -> Tuple:
-        field_value = self.search_widget.text()
-        if field_value == "":
-            field_value = None
-        return self.search_widget.objectName(), field_value
-
-
-class SearchResults(QTableView):
-
-    def __init__(self):
-        super().__init__()
-        self.doubleClicked.connect(lambda x: BasicSample.query(submitter_id=x.sibling(x.row(), 0).data()).show_details(self))
-
-    def setData(self, df:DataFrame) -> None:
-        """
-        sets data in model
-        """
-        self.data = df
-        try:
-            self.data['id'] = self.data['id'].apply(str)
-            self.data['id'] = self.data['id'].str.zfill(3)
-        except (TypeError, KeyError):
-            logger.error("Couldn't format id string.")
-        proxy_model = QSortFilterProxyModel()
-        proxy_model.setSourceModel(pandasModel(self.data))
-        self.setModel(proxy_model)
-        
\ No newline at end of file
diff --git a/src/submissions/frontend/widgets/submission_details.py b/src/submissions/frontend/widgets/submission_details.py
index a9b5f9a..c7af9fc 100644
--- a/src/submissions/frontend/widgets/submission_details.py
+++ b/src/submissions/frontend/widgets/submission_details.py
@@ -45,7 +45,8 @@ class SubmissionDetails(QDialog):
         self.btn.clicked.connect(self.export)
         self.back = QPushButton("Back")
         self.back.setFixedWidth(100)
-        self.back.clicked.connect(self.back_function)
+        # self.back.clicked.connect(self.back_function)
+        self.back.clicked.connect(self.webview.back)
         self.layout.addWidget(self.back, 0, 0, 1, 1)
         self.layout.addWidget(self.btn, 0, 1, 1, 9)
         self.layout.addWidget(self.webview, 1, 0, 10, 10)
@@ -63,8 +64,8 @@ class SubmissionDetails(QDialog):
                 self.reagent_details(reagent=sub)
         self.webview.page().setWebChannel(self.channel)
 
-    def back_function(self):
-        self.webview.back()
+    # def back_function(self):
+    #     self.webview.back()
 
     def activate_export(self):
         title = self.webview.title()
@@ -75,7 +76,11 @@ class SubmissionDetails(QDialog):
             # logger.debug(f"Updating export plate to: {self.export_plate}")
         else:
             self.btn.setEnabled(False)
-        if title == self.webview.history().items()[0].title():
+        try:
+            check = self.webview.history().items()[0].title()
+        except IndexError as e:
+            check = title
+        if title == check:
             # logger.debug("Disabling back button")
             self.back.setEnabled(False)
         else:
@@ -96,7 +101,7 @@ class SubmissionDetails(QDialog):
         exclude = ['submissions', 'excluded', 'colour', 'tooltip']
         base_dict['excluded'] = exclude
         template = sample.get_details_template()
-        template_path = Path(self.template.environment.loader.__getattribute__("searchpath")[0])
+        template_path = Path(template.environment.loader.__getattribute__("searchpath")[0])
         with open(template_path.joinpath("css", "styles.css"), "r") as f:
             css = f.read()
         html = template.render(sample=base_dict, css=css)
diff --git a/src/submissions/frontend/widgets/submission_widget.py b/src/submissions/frontend/widgets/submission_widget.py
index 26c9981..46ca734 100644
--- a/src/submissions/frontend/widgets/submission_widget.py
+++ b/src/submissions/frontend/widgets/submission_widget.py
@@ -31,6 +31,7 @@ class MyQComboBox(QComboBox):
     """
     Custom combobox that disables wheel events until focussed on.
     """
+
     def __init__(self, scrollWidget=None, *args, **kwargs):
         super(MyQComboBox, self).__init__(*args, **kwargs)
         self.scrollWidget = scrollWidget
@@ -48,6 +49,7 @@ class MyQDateEdit(QDateEdit):
     """
     Custom date editor that disables wheel events until focussed on.
     """
+
     def __init__(self, scrollWidget=None, *args, **kwargs):
         super(MyQDateEdit, self).__init__(*args, **kwargs)
         self.scrollWidget = scrollWidget
@@ -150,7 +152,7 @@ class SubmissionFormContainer(QWidget):
 
     @report_result
     def add_reagent(self, reagent_lot: str | None = None, reagent_role: str | None = None, expiry: date | None = None,
-                    name: str | None = None) -> Tuple[PydReagent, Report]:
+                    name: str | None = None, kit: str | KitType | None = None) -> Tuple[PydReagent, Report]:
         """
         Action to create new reagent in DB.
 
@@ -167,7 +169,8 @@ class SubmissionFormContainer(QWidget):
         if isinstance(reagent_lot, bool):
             reagent_lot = ""
         # NOTE: create form
-        dlg = AddReagentForm(reagent_lot=reagent_lot, reagent_role=reagent_role, expiry=expiry, reagent_name=name)
+        dlg = AddReagentForm(reagent_lot=reagent_lot, reagent_role=reagent_role, expiry=expiry, reagent_name=name,
+                             kit=kit)
         if dlg.exec():
             # NOTE: extract form info
             info = dlg.parse_form()
@@ -228,7 +231,7 @@ class SubmissionFormWidget(QWidget):
         # self.scrape_reagents(self.pyd.extraction_kit)
         self.scrape_reagents(self.extraction_kit)
 
-    def create_widget(self, key: str, value: dict | PydReagent, submission_type: str | SubmissionType| None = None,
+    def create_widget(self, key: str, value: dict | PydReagent, submission_type: str | SubmissionType | None = None,
                       extraction_kit: str | None = None, sub_obj: BasicSubmission | None = None,
                       disable: bool = False) -> "self.InfoItem":
         """
@@ -506,7 +509,8 @@ class SubmissionFormWidget(QWidget):
                     return None, None
             return self.input.objectName(), dict(value=value, missing=self.missing)
 
-        def set_widget(self, parent: QWidget, key: str, value: dict, submission_type: str | SubmissionType | None = None,
+        def set_widget(self, parent: QWidget, key: str, value: dict,
+                       submission_type: str | SubmissionType | None = None,
                        sub_obj: BasicSubmission | None = None) -> QWidget:
             """
             Creates form widget
@@ -682,9 +686,11 @@ class SubmissionFormWidget(QWidget):
                                     message=f"Couldn't find reagent type {self.reagent.role}: {lot} in the database.\n\nWould you like to add it?")
                 if dlg.exec():
                     wanted_reagent = self.parent().parent().add_reagent(reagent_lot=lot,
-                                                                           reagent_role=self.reagent.role,
-                                                                           expiry=self.reagent.expiry,
-                                                                           name=self.reagent.name)
+                                                                        reagent_role=self.reagent.role,
+                                                                        expiry=self.reagent.expiry,
+                                                                        name=self.reagent.name,
+                                                                        kit=self.extraction_kit
+                                                                        )
                     return wanted_reagent, report
                 else:
                     # NOTE: In this case we will have an empty reagent and the submission will fail kit integrity check
@@ -772,4 +778,3 @@ class SubmissionFormWidget(QWidget):
                 self.setObjectName(f"lot_{reagent.role}")
                 self.addItems(relevant_reagents)
                 self.setToolTip(f"Enter lot number for the reagent used for {reagent.role}")
-                
\ No newline at end of file
diff --git a/src/submissions/tools/__init__.py b/src/submissions/tools/__init__.py
index 2eba114..497a813 100644
--- a/src/submissions/tools/__init__.py
+++ b/src/submissions/tools/__init__.py
@@ -13,6 +13,7 @@ from jinja2 import Environment, FileSystemLoader
 from logging import handlers
 from pathlib import Path
 from sqlalchemy.orm import Session
+from sqlalchemy.orm.state import InstanceState
 from sqlalchemy import create_engine, text, MetaData
 from pydantic import field_validator, BaseModel, Field
 from pydantic_settings import BaseSettings, SettingsConfigDict
@@ -952,6 +953,10 @@ def report_result(func):
         match output:
             case Report():
                 report = output
+            # case InstanceState():
+            #     for attr in output.attrs:
+            #         print(f"{attr}: {attr.load_history()}")
+            #     return
             case tuple():
                 try:
                     report = [item for item in output if isinstance(item, Report)][0]
@@ -959,6 +964,7 @@ def report_result(func):
                     report = None
             case _:
                 report = None
+                return report
         logger.debug(f"Got report: {report}")
         try:
             results = report.results
@@ -982,3 +988,5 @@ def report_result(func):
         # logger.debug(f"Returning true output: {true_output}")
         return true_output
     return wrapper
+
+