#!/bin/bash

function print_help_and_exit
{
cat >&2 << 'EOF'

'voronota-js-iface-visualize' script generates inter-chain interface visualizations.

Options:
    --input | -i              string  *  input file path
    --output | -o             string  *  output directory path
    --restrict-input          string     query to restrict input atoms, default is '[-t! het]'
    --select-contacts         string     query to select contacts, default is '[-inter-chain]'
    --subselect-contacts      string     query to subselect contacts, default is '[]'
    --as-assembly                        flag to treat input file as biological assembly
    --help | -h                          flag to display help message and exit

Standard output:
    list of generated files
    
Examples:

    voronota-js-iface-visualize --input model1.pdb --output results

EOF
exit 1
}

readonly ZEROARG=$0
ALLARGS=("$@")

if [[ $ZEROARG == *"/"* ]]
then
	cd "$(dirname ${ZEROARG})"
	export PATH="$(pwd):${PATH}"
	cd - &> /dev/null
fi

export LC_ALL=C

command -v voronota-js &> /dev/null || { echo >&2 "Error: 'voronota-js' executable not in binaries path"; exit 1; }
command -v voronota &> /dev/null || { echo >&2 "Error: 'voronota' executable not in binaries path"; exit 1; }

INFILE=""
OUTDIR=""
RESTRICT_INPUT="[-t! het]"
SELECT_CONTACTS="[-inter-chain]"
SUBSELECT_CONTACTS="[]"
AS_ASSEMBLY="false"
HELP_MODE="false"

while [[ $# > 0 ]]
do
	OPTION="$1"
	OPTARG="$2"
	shift
	case $OPTION in
	-i|--input)
		INFILE="$OPTARG"
		shift
		;;
	-o|--output)
		OUTDIR="$OPTARG"
		shift
		;;
	--restrict-input)
		RESTRICT_INPUT="$OPTARG"
		shift
		;;
	--select-contacts)
		SELECT_CONTACTS="$OPTARG"
		shift
		;;
	--subselect-contacts)
		SUBSELECT_CONTACTS="$OPTARG"
		shift
		;;
	--as-assembly)
		AS_ASSEMBLY="true"
		;;
	-h|--help)
		HELP_MODE="true"
		;;
	*)
		echo >&2 "Error: invalid command line option '$OPTION'"
		exit 1
		;;
	esac
done

if [ -z "$INFILE" ] || [ -z "$OUTDIR" ] || [ "$HELP_MODE" == "true" ]
then
	print_help_and_exit
fi

if [ ! -s "$INFILE" ]
then
	echo >&2 "Error: input file '$INFILE' does not exist"
	exit 1
fi

mkdir -p "$OUTDIR"

if [ ! -d "$OUTDIR" ]
then
	echo >&2 "Error: failed to create output directory '$OUTDIR'"
	exit 1
fi

readonly TMPLDIR=$(mktemp -d)
trap "rm -r $TMPLDIR" EXIT

{
cat << EOF

params={}

params.input_file='$INFILE';
params.input_as_assembly='$AS_ASSEMBLY';
params.restrict_input_atoms='$RESTRICT_INPUT';
params.contacts_selection='$SELECT_CONTACTS';
params.contacts_subselection='$SUBSELECT_CONTACTS';
params.output_prefix='${TMPLDIR}/output';

EOF

cat << 'EOF'

if(params.input_file===undefined || params.input_file==="")
{
	throw ("No input file");
}

if(params.input_as_assembly===undefined || params.input_as_assembly==="")
{
	params.input_as_assembly="false";
}

if(params.restrict_input_atoms===undefined || params.restrict_input_atoms==="")
{
	params.restrict_input_atoms='[]';
}

if(params.contacts_subselection===undefined || params.contacts_subselection==="")
{
	params.contacts_subselection='[]';
}

voronota_import("-file", params.input_file, "-as-assembly", params.input_as_assembly, "-include-heteroatoms");
voronota_assert_partial_success("Failed to import file");

voronota_restrict_atoms("-use", params.restrict_input_atoms);
voronota_assert_full_success("Failed to restrict input atoms by the input query");

voronota_export_atoms("-as-pdb", "-file", params.output_prefix+"_atoms.pdb");
voronota_assert_full_success("Failed to export atoms");

voronota_construct_contacts("-calculate-bounding-arcs -calculate-adjacencies");
voronota_assert_full_success("Failed to construct contacts");

voronota_select_contacts("-use", "("+params.contacts_selection+" and "+params.contacts_subselection+")", "-name", "inter_chain_contacts")
voronota_assert_full_success("Failed to select inter-chain contacts");

voronota_export_adjuncts_of_contacts("-file", params.output_prefix+"_contacts_graph_vertices.txt", "-adjacency-file", params.output_prefix+"_contacts_graph_edges.txt", "-contacts-use [inter_chain_contacts] -no-serial",
  "-adjuncts contact_index ir_contact_index area distance boundary");
voronota_assert_full_success("Failed to export contacts graph");

voronota_export_adjuncts_of_contacts("-file", params.output_prefix+"_ir_contacts_graph_vertices.txt", "-adjacency-file", params.output_prefix+"_ir_contacts_graph_edges.txt", "-contacts-use [inter_chain_contacts] -no-serial",
  "-adjuncts contact_index ir_contact_index area distance boundary -inter-residue");
voronota_assert_full_success("Failed to export contacts graph");


voronota_select_atoms("-use", "[-sel-of-contacts inter_chain_contacts]", "-name", "inter_chain_contacts_atoms")
voronota_assert_full_success("Failed to select inter-chain contacts atoms");

voronota_export_adjuncts_of_atoms("-file", params.output_prefix+"_atoms_graph_vertices.txt", "-use [inter_chain_contacts_atoms] -no-serial -adjuncts atom_index residue_index");
voronota_assert_full_success("Failed to export atoms adjuncts file");

voronota_export_adjuncts_of_contacts("-file", params.output_prefix+"_atoms_graph_edges.txt", "-atoms-use [inter_chain_contacts_atoms] -contacts-use [-no-solvent] -no-serial -adjuncts atom_index1 atom_index2 area");
voronota_assert_full_success("Failed to export contacts adjuncts file");

voronota_export_adjuncts_of_contacts("-file", params.output_prefix+"_atoms_graph_edges_interchain.txt", "-atoms-use [inter_chain_contacts_atoms] -contacts-use [inter_chain_contacts] -no-serial -adjuncts atom_index1 atom_index2 area");
voronota_assert_full_success("Failed to export contacts adjuncts file");

voronota_make_drawable_contacts("-use [inter_chain_contacts]");

voronota_color_contacts("-col green");
voronota_export_contacts_as_pymol_cgo("-use [inter_chain_contacts] -wireframe -name wireframe", "-file", params.output_prefix+"_draw_contacts_wireframe.py");

voronota_color_contacts("-col red");
voronota_export_contacts_as_pymol_cgo("-use [inter_chain_contacts] -fat-wireframe -fat-wireframe-min-radius 1.2 -fat-wireframe-max-radius 1.39 -name boundary_tangent_spheres",
  "-file", params.output_prefix+"_draw_contacts_boundary_tangent_spheres.py");
  
voronota_color_contacts("-col yellow");
voronota_export_contacts_as_pymol_cgo("-use [inter_chain_contacts] -fat-wireframe -fat-wireframe-min-radius 0.0 -fat-wireframe-max-radius 1.2 -name edges_tangent_spheres",
  "-file", params.output_prefix+"_draw_contacts_edges_tangent_spheres.py");
  
voronota_color_contacts("-col magenta");
voronota_export_contacts_as_pymol_cgo("-use [inter_chain_contacts] -fat-wireframe -fat-wireframe-peaks-only -fat-wireframe-min-radius 0.0 -fat-wireframe-max-radius 1.39 -name peak_tangent_spheres",
  "-file", params.output_prefix+"_draw_contacts_peak_tangent_spheres.py");
  
voronota_color_contacts("-col white");
voronota_export_contacts_as_pymol_cgo("-use [inter_chain_contacts] -fat-wireframe -fat-wireframe-peaks-only -fat-wireframe-min-radius 0.0 -fat-wireframe-max-radius 1.39 -fat-wireframe-transform-a 0.0 -fat-wireframe-transform-b 0.2 -name peak_tangent_sphere_centers",
  "-file", params.output_prefix+"_draw_contacts_peak_tangent_sphere_centers.py");

EOF

} \
| voronota-js


cd "$TMPLDIR"

{
echo "graph G {"
cat output_contacts_graph_vertices.txt | tail -n +2 | voronota expand-descriptors \
| awk '{print ($15+1) " [" "cluster=\"" ($16+1) "\", " "label=\"" $6 "." $7 "\\n" $13 "." $14 "\"," "fontsize=\"" ($17*1.2+7.0) "\"," "margin=0,shape=\"circle\"" "]"}'
cat output_contacts_graph_edges.txt | tail -n +2 \
| awk '{print ($1+1) " -- " ($2+1)}'
echo "}"
} \
> ./output_contacts_graph.dot

sfdp -Goverlap="prism" ./output_contacts_graph.dot -Tpng > output_contacts_graph.png
sfdp -Goverlap="prism" ./output_contacts_graph.dot | gvmap  -e | neato -Ecolor="#00000055" -Gbgcolor="white" -n10 -Tpng > output_contacts_graph_gvmap.png

{
echo "graph G {"
cat output_ir_contacts_graph_vertices.txt | tail -n +2 | voronota expand-descriptors \
| awk '{print ($16+1) " [" "cluster=\"" ($16+1) "\", " "label=\"" $6 "\\n" $13 "\"," "fontsize=\"" ($17*1.2+7.0) "\"," "margin=0,shape=\"circle\"" "]"}'
cat output_ir_contacts_graph_edges.txt | tail -n +2 \
| awk '{print ($1+1) " -- " ($2+1)}'
echo "}"
} \
> ./output_ir_contacts_graph.dot

sfdp -Goverlap="prism" ./output_ir_contacts_graph.dot -Tpng > output_ir_contacts_graph.png
sfdp -Goverlap="prism" ./output_ir_contacts_graph.dot | gvmap  -e | neato -Ecolor="#00000055" -Gbgcolor="white" -n10 -Tpng > output_ir_contacts_graph_gvmap.png

{
echo "graph G {"
cat output_atoms_graph_vertices.txt | tail -n +2 | voronota expand-descriptors \
| awk '{print ($8+1) " [" "cluster=\"" ($9+1) "\", " "label=\"" $1 "." $6 "." $7 "\"," "margin=0,shape=\"circle\"" "]"}'
cat output_atoms_graph_edges.txt | tail -n +2 \
| awk '{print ($3+1) " -- " ($4+1)}'
echo "}"
} \
> ./output_atoms_graph.dot

sfdp -Goverlap="prism" ./output_atoms_graph.dot -Tpng > output_atoms_graph.png
sfdp -Goverlap="prism" ./output_atoms_graph.dot | gvmap  -e | neato -Ecolor="#00000055" -Gbgcolor="white" -n10 -Tpng > output_atoms_graph_gvmap.png

{
echo "graph G {"
cat output_atoms_graph_vertices.txt | tail -n +2 | voronota expand-descriptors \
| awk '{print ($8+1) " [" "cluster=\"" ($9+1) "\", " "label=\"" $1 "." $6 "." $7 "\"," "margin=0,shape=\"circle\"" "]"}'
cat output_atoms_graph_edges_interchain.txt | tail -n +2 \
| awk '{print ($3+1) " -- " ($4+1)}'
echo "}"
} \
> ./output_atoms_graph_interchain.dot

sfdp -Goverlap="prism" ./output_atoms_graph_interchain.dot -Tpng > output_atoms_graph_interchain.png
sfdp -Goverlap="prism" ./output_atoms_graph_interchain.dot | gvmap  -e | neato -Ecolor="#00000055" -Gbgcolor="white" -n10 -Tpng > output_atoms_graph_interchain_gvmap.png

cd - &> /dev/null

mv $TMPLDIR/* "$OUTDIR/"

