"""
Simple distogram plotting
=============================

This script contains a simple example of how you can plot
distograms using ConKit

"""

import conkit.io
import conkit.plot

# Define the input variables
sequence_file = "7l6q/7l6q.fasta"
sequence_format = "fasta"
distance_file = "7l6q/7l6q.af2"
distance_format = "alphafold2"

# Create ConKit hierarchies
#       Note, we only need the first Sequence/Distogram
#       from each file
seq = conkit.io.read(sequence_file, sequence_format).top
distpred = conkit.io.read(distance_file, distance_format).top

# Assign the sequence register to your contact prediction
distpred.sequence = seq
distpred.set_sequence_register()

# Then we can plot the distogram predicted with AlphaFold 2
fig = conkit.plot.DistogramHeatmapFigure(distpred)
fig.savefig("7l6q/7l6q.png")