From ccb43aafde11c2640d665104b1c62539aa34bf7f Mon Sep 17 00:00:00 2001
From: sverhoeven <s.verhoeven@esciencecenter.nl>
Date: Thu, 28 Mar 2024 17:13:51 +0100
Subject: [PATCH] Add colors to sequence + started on chain selection in 3D

---
 app/scenarios/MoleculeSubForm.client.tsx |  5 ++++-
 app/scenarios/ResiduesSelect.tsx         | 26 +++++++++++++++++++---
 app/scenarios/Viewer.client.tsx          | 28 ++++++++++++++++++------
 app/scenarios/molecule.client.ts         | 21 +++++++++++++++++-
 4 files changed, 68 insertions(+), 12 deletions(-)

diff --git a/app/scenarios/MoleculeSubForm.client.tsx b/app/scenarios/MoleculeSubForm.client.tsx
index e2819ce4..cdd551d4 100644
--- a/app/scenarios/MoleculeSubForm.client.tsx
+++ b/app/scenarios/MoleculeSubForm.client.tsx
@@ -98,7 +98,10 @@ export function MoleculeSubForm({
       </FormItem>
       <div className="h-[500px] w-full">
         {molecule ? (
-          <Viewer structure={molecule?.structure} />
+          <Viewer
+            structure={molecule?.structure}
+            chain={actpass.active.chain}
+          />
         ) : (
           <p>Load a structure first</p>
         )}
diff --git a/app/scenarios/ResiduesSelect.tsx b/app/scenarios/ResiduesSelect.tsx
index c2c33cf8..e973b859 100644
--- a/app/scenarios/ResiduesSelect.tsx
+++ b/app/scenarios/ResiduesSelect.tsx
@@ -1,6 +1,14 @@
 import { ChangeEvent, useState } from "react";
 import { FormDescription } from "./FormDescription";
-import { Residue } from "./molecule.client";
+import { Residue, SecondaryStructure } from "./molecule.client";
+import clsx from "clsx";
+
+const residueVariants: Record<SecondaryStructure, string> = {
+  sheet: "bg-amber-200 dark:bg-amber-700",
+  helix: "bg-red-200 dark:bg-red-900",
+  turn: "",
+  "": "",
+};
 
 export function ResiduesSelect({
   options,
@@ -46,13 +54,19 @@ export function ResiduesSelect({
           return (
             <div
               key={r.resno}
-              className="inline-block w-4"
+              className="inline-block w-4 text-center font-mono"
               title={`${r.resno}`}
             >
-              <span>{r.seq}</span>
+              <label
+                htmlFor={`residue-${r.resno}`}
+                className={residueVariants[r.sec]}
+              >
+                {r.seq}
+              </label>
               <input
                 type="checkbox"
                 value={r.resno}
+                id={`residue-${r.resno}`}
                 checked={selected.includes(r.resno)}
                 onChange={(e) => handleChange(e, index)}
               />
@@ -61,6 +75,12 @@ export function ResiduesSelect({
         })}
       </div>
       <FormDescription>
+        <span className={clsx("p-1 font-mono", residueVariants["helix"])}>
+          Helix
+        </span>{" "}
+        <span className={clsx("p-1 font-mono", residueVariants["sheet"])}>
+          Sheet
+        </span>{" "}
         (Hold Shift to select a range of residues)
       </FormDescription>
       {/* TODO add buttons to select all, none, invert */}
diff --git a/app/scenarios/Viewer.client.tsx b/app/scenarios/Viewer.client.tsx
index 3360c06a..b1a4733b 100644
--- a/app/scenarios/Viewer.client.tsx
+++ b/app/scenarios/Viewer.client.tsx
@@ -1,7 +1,13 @@
 import { useEffect, useRef, useState } from "react";
-import { Stage, Structure } from "ngl";
+import { Stage, Structure, StructureComponent } from "ngl";
 
-export function Viewer({ structure }: { structure: Structure }) {
+export function Viewer({
+  structure,
+  chain,
+}: {
+  structure: Structure;
+  chain?: string;
+}) {
   const viewportRef = useRef<HTMLDivElement>(null);
   const stage = useRef<Stage | null>(null);
   const [isLoaded, setIsLoaded] = useState(false);
@@ -45,11 +51,7 @@ export function Viewer({ structure }: { structure: Structure }) {
       console.error("Could not load structure");
       return;
     }
-    component.addRepresentation("cartoon", { sele: "polymer" });
-    component.addRepresentation("ball+stick", { sele: "ligand" });
-    component.addRepresentation("base", { sele: "nucleic" });
-
-    component.autoView();
+    stage.current.defaultFileRepresentation(component);
 
     setIsLoaded(true);
 
@@ -65,5 +67,17 @@ export function Viewer({ structure }: { structure: Structure }) {
     // };
   }, [structure, isLoaded]);
 
+  useEffect(() => {
+    if (stage.current === null) {
+      return;
+    }
+
+    stage.current.eachRepresentation((repr) => {
+      const selection = chain ? `:${chain}` : "";
+      // TODO figure out how to set selection
+      // repr.setSelection(selection);
+    });
+  }, [chain]);
+
   return <div ref={viewportRef} className="h-full w-full"></div>;
 }
diff --git a/app/scenarios/molecule.client.ts b/app/scenarios/molecule.client.ts
index c62d4c7d..52773376 100644
--- a/app/scenarios/molecule.client.ts
+++ b/app/scenarios/molecule.client.ts
@@ -1,9 +1,27 @@
 import { Structure, autoLoad } from "ngl";
+import ResidueProxy from "ngl/dist/declarations/proxy/residue-proxy";
 
-export type Residue = { resno: number; seq: string };
+export type Residue = { resno: number; seq: string; sec: SecondaryStructure };
 export type Chains = Record<string, Residue[]>;
 export type Molecule = { structure: Structure; chains: Chains; file: File };
 
+export type SecondaryStructure = "sheet" | "helix" | "turn" | "";
+
+function secondaryStructureOfResidue(
+  residue: ResidueProxy
+): SecondaryStructure {
+  if (residue.isSheet()) {
+    return "sheet";
+  }
+  if (residue.isHelix()) {
+    return "helix";
+  }
+  if (residue.isTurn()) {
+    return "turn";
+  }
+  return "";
+}
+
 export function chainsFromStructure(structure: Structure) {
   const chains: Chains = {};
   structure.eachChain((c) => {
@@ -13,6 +31,7 @@ export function chainsFromStructure(structure: Structure) {
       residues.push({
         resno: r.resno,
         seq: r.getResname1(),
+        sec: secondaryStructureOfResidue(r),
       });
     });
     chains[chainName] = residues;