#!/bin/bash

# Check if directory name is provided
if [ -z "$1" ]; then
    echo "Usage: $0 <directory_name>"
    exit 1
fi

DIR_NAME="$1"

# Check if directory exists
if [ -d "$DIR_NAME" ]; then
    echo "Error: Directory '$DIR_NAME' already exists. Aborting."
    exit 1
fi

# Create main and subdirectories
mkdir -p "$DIR_NAME/af_input" "$DIR_NAME/af_output" "$DIR_NAME/af_weights" "$DIR_NAME/af_public_databases" "$DIR_NAME/tmp" 

cat <<EOL > "$DIR_NAME/af_input/TEMPLATE_model.json"

{
  "name": "Hello fold",
  "modelSeeds": [10, 42],
  "sequences": [
    {
      "protein": {
        "id": "A",
        "sequence": "PVLSCGEWQL",
        "modifications": [
          {"ptmType": "HY3", "ptmPosition": 1},
          {"ptmType": "P1L", "ptmPosition": 5}
        ],
        "unpairedMsa": ...,
      }
    },
    {
      "protein": {
        "id": "B",
        "sequence": "RPACQLW",
        "templates": [
          {
            "mmcif": ...,
            "queryIndices": [0, 1, 2, 4, 5, 6],
            "templateIndices": [0, 1, 2, 3, 4, 8]
          }
        ]
      }
    },
    {
      "dna": {
        "id": "C",
        "sequence": "GACCTCT",
        "modifications": [
          {"modificationType": "6OG", "basePosition": 1},
          {"modificationType": "6MA", "basePosition": 2}
        ]
      }
    },
    {
      "rna": {
        "id": "E",
        "sequence": "AGCU",
        "modifications": [
          {"modificationType": "2MG", "basePosition": 1},
          {"modificationType": "5MC", "basePosition": 4}
        ],
        "unpairedMsa": ...
      }
    },
    {
      "ligand": {
        "id": ["F", "G", "H"],
        "ccdCodes": ["ATP"]
      }
    },
    {
      "ligand": {
        "id": "I",
        "ccdCodes": ["NAG", "FUC"]
      }
    },
    {
      "ligand": {
        "id": "Z",
        "smiles": "CC(=O)OC1C[NH+]2CCC1CC2"
      }
    }
  ],
  "bondedAtomPairs": [
    [["A", 1, "CA"], ["B", 1, "CA"]],
    [["A", 1, "CA"], ["G", 1, "CHA"]],
    [["J", 1, "O6"], ["J", 2, "C1"]]
  ],
  "userCcd": ...,
  "dialect": "alphafold3",
  "version": 1
}


EOL
echo "Created template model  variable in $DIR_NAME/af_input/TEMPLATE_model.json"


############################################################################################
##################### Create AF_ENV file with example environment variables ################
############################################################################################

cat <<EOL > "$DIR_NAME/AF_ENV"
#!/bin/bash
# AF_ENV - Environment Variables for AlphaFold Job

# model weights are subjected to license approval, and you should obtain them independently from the HPC team. 
# Edit the following path to point towards the model weights 
export WEIGHT_DIR=/path/to/data
# Alphafold will need a JSON file to be executed. define here the name of the JSON file.
# ... and copy it to the input folder
export JSON_NAME=testrun.json

# do not edit after this line
export AF_CONTAINER=/scratch/container/alphafold_noArch-3.0.1.sif
export OUTPUT_DIR=\$(pwd)/af_output
export INPUT_DIR=\$(pwd)/af_input
export TMP_DIR=\$(pwd)/tmp
export DBASE_DIR=/scratch/data/alphafold/3.0.1 
EOL
echo "Created environment variable in $DIR_NAME/AF_ENV"
# 

############################################################################################
############################ Create SLURM job script template ##############################
############################################################################################


cat <<EOL > "$DIR_NAME/TEMPLATE_slurm.sh"
#!/bin/bash
#SBATCH --job-name=TestAlphaFoldContainer
#SBATCH --output=af_output/slurm-%j.out
#SBATCH --error=af_output/slurm-%j.err
#SBATCH --time=2:00:00
#SBATCH --partition=gpuhgx 
#SBATCH --gres=gpu:1
#SBATCH --cpus-per-task=8
#SBATCH --mem=128G
#SBATCH --mail-type=ALL
#SBATCH --mail-user=$USER@uni-muenster.de

# Load environment
source ./AF_ENV

# module loads
ml Apptainer

# Run AlphaFold job from the system-wide container
apptainer exec \\
	--containall \\
	--nv \\
	--mount type=bind,src=\$(pwd)/,dst=\$HOME \\
	--mount type=bind,src=\$DBASE_DIR/,dst=\$HOME/af_public_databases \\
	--mount type=bind,src=\$WEIGHT_DIR/,dst=\$HOME/af_weights \\
	--mount type=bind,src=\$TMP_DIR/,dst=/tmp \\
	\$AF_CONTAINER \\
	bash -c "python /app/alphafold/run_alphafold.py --json_path=./af_input/\$JSON_NAME --model_dir=./af_weights  --db_dir=./af_public_databases  --output_dir=./af_output"

EOL

echo "Created sbatch template file in $DIR_NAME/submit_job.TEMPLATE"
echo
echo "project folder summary:"
find $DIR_NAME -print | sort|sed -e 's;[^/]*/;|____;g;s;____|; |;g'
echo
sleep 3
echo "Setup completed in directory '$DIR_NAME'"