import argparse

from biopandas.pdb import PandasPdb


def get_cli_arguments():
    """Command line argument parser.

    Returns
    -------
    argparse.Namespace
        Object containing arguments.
    """
    parser = argparse.ArgumentParser(add_help=False)
    required = parser.add_argument_group("required arguments")
    optional = parser.add_argument_group("optional arguments")
    required.add_argument(
        "--input-pdb",
        action="store",
        type=str,
        help="Input PDB file containing the 3D structure of the genome.",
        required=True,
    )
    required.add_argument(
        "--output-pdb",
        action="store",
        type=str,
        help="Output PDB file containing the completed 3D structure of the genome.",
        required=True,
    )
    # Add help.
    optional.add_argument(
        "-h",
        "--help",
        action="help",
        default=argparse.SUPPRESS,
        help="Show this help message and exit.",
    )
    return parser.parse_args()


def interpolate_missing_coordinates(pdb_name_in, pdb_name_out):
    """Interpolate missing coordinates in PDB file containing a 3D genome structure.

    3D coordinates of the missing beads are interpolated one by one.

    Parameters
    ----------
    pdb_name_in : str
        PDB file containing the 3D structure of the genome
    pdb_name_out : str
        Output PDB file containing the annotated 3D structure of the genome
    """
    pdb = PandasPdb().read_pdb(pdb_name_in)
    atoms = pdb.df["ATOM"]
    print(f"Found {atoms.shape[0]} beads in {pdb_name_in}")
    # Interpolate missing coordinates chromosome by chromosome.
    coord_columns= ["x_coord", "y_coord", "z_coord"]
    for residue_number in atoms["residue_number"].unique():
        missing_beads = (
            atoms.query(f"residue_number == {residue_number}")
            [coord_columns]
            .isna()
            .sum()   # Sum NAs over columns (coordinates).
            .sum()   # Sum NAs over lines.
        ) // 3       # 1     missing atom has 3 missing coordinates.
        chromosome = atoms.query(f"residue_number == {residue_number}")
        if not missing_beads:
            continue
        print(f"Chromosome {residue_number} has {missing_beads} missing beads")

        # Interpolate coordinates on the three coordinates (x, y, z).
        # Only interpolation between existing beads.
        # No extrapolation at chromosome ends.
        interpolate_chromosome = chromosome.copy()
        interpolate_chromosome[coord_columns] = chromosome[coord_columns].interpolate(
            method="pchip", axis=0, limit_area="inside"
        )
        # Print number of missing beads that will later remove.
        deleted_atoms_number = interpolate_chromosome["x_coord"].isna().sum()
        print(
            f"Removed {deleted_atoms_number} beads from chromosome {residue_number} extremities"
        )
        # Save interpolated chromosome into the full structure.
        atoms.loc[atoms["residue_number"] == residue_number, coord_columns] = interpolate_chromosome[coord_columns]
    # Remove in one step all beads with missing coordinates
    # located at chromosome extremities.
    atoms = atoms.dropna(subset=coord_columns)
    print(
        f"Found {atoms.shape[0]} beads after interpolation and removal of extremities missing beads."
    )

    pdb.df["ATOM"] = atoms
    pdb.to_pdb(pdb_name_out)
    print(f"Wrote {pdb_name_out}")


if __name__ == "__main__":
    ARGS = get_cli_arguments()

    interpolate_missing_coordinates(ARGS.input_pdb, ARGS.output_pdb)

import argparse
import math
from pathlib import Path

from Bio import SeqIO
from biopandas.pdb import PandasPdb


def is_file(parser, file_path):
    """Check file exists.

    Parameters
    ----------
    parser : argparse.ArgumentParser
        Command line argument parser.
    file_path : str
        File path
    Returns
    -------
    str
        File path.
    """
    if not Path(file_path).is_file():
        parser.error(f"The file {file_path} does not exist")
    else:
        return file_path


def get_cli_arguments():
    """Command line argument parser.

    Returns
    -------
    argparse.Namespace
        Object containing arguments.
    """
    parser = argparse.ArgumentParser(add_help=False)
    required = parser.add_argument_group("required arguments")
    optional = parser.add_argument_group("optional arguments")
    required.add_argument(
        "--pdb",
        action="store",
        type=lambda name: is_file(parser, name),
        help="PDB file containing the 3D structure of the genome.",
        required=True,
    )
    required.add_argument(
        "--fasta",
        action="store",
        type=lambda name: is_file(parser, name),
        help="FASTA file containing the sequence of the genome.",
        required=True,
    )
    required.add_argument(
        "--resolution",
        action="store",
        type=int,
        help="HiC resolution from which the structure has been generated.",
        required=True,
    )
    required.add_argument(
        "--output",
        action="store",
        type=str,
        help="Output PDB file containing the annotated 3D structure of the genome.",
        required=True,
    )
    # Add help.
    optional.add_argument(
        "-h",
        "--help",
        action="help",
        default=argparse.SUPPRESS,
        help="Show this help message and exit.",
    )
    return parser.parse_args()


def extract_chromosome_length(fasta_name):
    """Extract chromosome length from a FASTA file.

    Parameters
    ----------
    fasta_name : str
        Name of Fasta file containing the sequence of the genome.

    Returns
    -------
    list
        List of chromosome lengths.
    """
    chromosome_length_lst = []
    with open(fasta_name, "r") as fasta_file:
        print(f"Reading {fasta_name}")
        for record in SeqIO.parse(fasta_file, "fasta"):
            length = len(record.seq)
            print(f"Found chromosome {record.id} with {length} bases")
            chromosome_length_lst.append(length)
    return chromosome_length_lst


def assign_chromosome_number(
    pdb_name_in, chromosome_length, HiC_resolution, pdb_name_out
):
    """Assign chromosome numbers to the whole genome 3D structure.

    Note:
    - The PDB file produced by Pastis is not readable by Biopython
    because the residue number column is missing.
    - We use instead the biopandas library: http://rasbt.github.io/biopandas/

    Parameters
    ----------
    pdb_name_in : str
        PDB file containing the 3D structure of the genome.
    chromosome_length : list
        List with chromosome lengths.
    HiC_resolution : int
        HiC resolution.
    pdb_name_out : str
        Output PDB file containing the annotated 3D structure of the genome.
    """
    coordinates = PandasPdb().read_pdb(pdb_name_in)
    coordinates_df = coordinates.df["ATOM"]
    # Verify the number of beads in the structure
    # and the number of beads deduced from the sequence and resolution
    # are the same.
    print(f"Number of beads read from structure: {coordinates_df.shape[0]}")
    beads_per_chromosome = [
        math.ceil(length / HiC_resolution) for length in chromosome_length
    ]
    print(
        f"Number of beads deduced from sequence and HiC resolution: {sum(beads_per_chromosome)}"
    )
    if coordinates_df.shape[0] != sum(beads_per_chromosome):
        raise ValueError(
            "Number of beads in structure and "
            "number of beads deduced from sequence and resolution "
            "are different."
        )
    # Add residue (chromosome) number.
    residue_number = []
    for index, bead_number in enumerate(beads_per_chromosome):
        residue_number = residue_number + [index + 1] * bead_number
    coordinates_df["residue_number"] = residue_number

    # Add residue name based on residue number.
    coordinates_df["residue_name"] = coordinates_df["residue_number"].apply(
        lambda x: f"C{x:02d}"
    )

    # Add chain (chromosome) name.
    # Chromosome 1 -> Chain A
    # Chromosome 2 -> Chain B
    # ...
    # Chromosome 26 -> Chain Z
    # Chromosome 27 -> Chain A
    # Chromosome 28 -> Chain B
    # ...
    letters = "ABCDEFGHIJKLMNOPQRSTUVWXYZ"
    # Handle case where there are more than 26 chromosomes.
    repeat_factor = math.ceil(len(chromosome_length) / len(letters))
    residue_to_chain = dict(
        zip(coordinates_df["residue_number"].unique(), letters * repeat_factor)
    )
    # Map residue number to chain id.
    coordinates_df["chain_id"] = coordinates_df["residue_number"].map(residue_to_chain)

    # Fix atom number starting at 1 (instead of 0).
    coordinates_df["atom_number"] = coordinates_df["atom_number"] + 1

    # Write new PDB file to disk.
    coordinates.df["ATOM"] = coordinates_df
    coordinates.to_pdb(path=pdb_name_out, records=None, gz=False, append_newline=True)
    print(f"Wrote {pdb_name_out}")


if __name__ == "__main__":
    ARGS = get_cli_arguments()

    # Read FASTA file and get chromosome length.
    CHROMOSOME_LENGTH = extract_chromosome_length(ARGS.fasta)

    # Assign chromosome number.
    assign_chromosome_number(ARGS.pdb, CHROMOSOME_LENGTH, ARGS.resolution, ARGS.output)

import argparse

import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
import numpy as np


def get_cli_arguments():
    """Parse command line.

    Returns
    -------
    argparse.Namespace
        Object containing arguments
    """
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--contacts",
        action="store",
        type=str,
        help="Input contact file",
        required=True,
    )
    parser.add_argument(
        "--map",
        action="store",
        type=str,
        help="Output contact map image",
        required=True,
    )
    return parser.parse_args()


def create_contact_map(contacts_file, map_file):
    """Create contact map image.

    Parameters
    ----------
    contacts_file : str
        Contact file
    map_file : str
        Contact map image
    """
    contacts = np.loadtxt(contacts_file)
    plt.imshow(np.log2(contacts+1), cmap="hot")
    plt.colorbar()
    plt.savefig(map_file, dpi=300)


if __name__ == "__main__":
    ARGS = get_cli_arguments()
    create_contact_map(ARGS.contacts, ARGS.map)

import argparse
import math
import pandas as pd
from pathlib import Path

from Bio import SeqIO
from biopandas.pdb import PandasPdb


def is_file(parser, file_path):
    """Check file exists.

    Parameters
    ----------
    parser : argparse.ArgumentParser
        Command line argument parser
    file_path : str
        File path
    Returns
    -------
    str
        File path    
    """
    if not Path(file_path).is_file():
        parser.error(f"The file {file_path} does not exist")
    else:
        return file_path


def get_cli_arguments():
    """Command line argument parser.

    Returns
    -------
    argparse.Namespace
        Object containing arguments
    """
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--pdb",
        action="store",
        type=lambda name: is_file(parser, name),
        help="PDB file containing the 3D structure of the genome",
        required=True,
    )
    parser.add_argument(
        "--fasta",
        action="store",
        type=lambda name: is_file(parser, name),
        help="Fasta file containing the sequence of the genome",
        required=True,
    )
    parser.add_argument(
        "--resolution", action="store", type=int, help="HiC resolution", required=True,
    )
    parser.add_argument(
        "--output",
        action="store",
        type=str,
        help="Output g3d file containing the annotated 3D structure of the genome",
        required=True,
    )
    return parser.parse_args()


def extract_chromosome_length(fasta_name):
    """Extract chromosome length from a FASTA file.

    Parameters
    ----------
    fasta_name : str
        Name of Fasta file containing the sequence of the genome

    Returns
    -------
    list
        List of chromosome lengthes
    """
    chromosome_length_lst = []
    with open(fasta_name, "r") as fasta_file:
        print(f"Reading {fasta_name}")
        for record in SeqIO.parse(fasta_file, "fasta"):
            length = len(record.seq)
            print(f"Found chromosome {record.id} with {length} bases")
            chromosome_length_lst.append(length)
    return chromosome_length_lst


def convert_to_g3d(pdb_name_in, chromosome_length, HiC_resolution, g3d_name_out):
    """Assign chromosome numbers to the whole genome 3D structure.

    Note:
    - The PDB file produced by Pastis is not readable by Biopython
    because the residue number column is missing.
    - We use instead the biopandas library. http://rasbt.github.io/biopandas/

    Parameters
    ----------
    pdb_name_in : str
        PDB file containing the 3D structure of the genome
    chromosome_length : list
        List with chromosome lengths
    HiC_resolution : int
        HiC resolution
    g3d_name_out : str
        Output g3d file containing the annotated 3D structure of the genome
    """
    pdb = PandasPdb().read_pdb(pdb_name_in)
    atoms = pdb.df["ATOM"]
    print(f"Number of beads read from structure: {pdb.df['ATOM'].shape[0]}")

    # Deduce chromosomal position of each atoms
    bp_coordinates = list()
    for i in range(len(chromosome_length)):
        bp_coordinates_chrom_x = list(range(0, chromosome_length[i], HiC_resolution))
        bp_coordinates = bp_coordinates + bp_coordinates_chrom_x

    # Get the index of the beads without the ones already removed from the pdb file (by interpolate_missing_coordinates.py )
    atom_number = list(atoms["atom_number"]-1)

    # Delete the bp coordinates values corresponding to removed beads
    bp_coordinates = pd.Series(bp_coordinates)
    bp_coordinates = bp_coordinates.iloc[atom_number]
    bp_coordinates.reset_index(drop=True, inplace=True)

    # Extract chromosomes and 3d coordinates information
    g3d_data = { 'chrom': atoms["residue_number"],
                'locus':bp_coordinates,
                'X3D_x': atoms["x_coord"],
                'X3D_y': atoms["y_coord"],
                'X3D_z': atoms["z_coord"] }
    result = pd.DataFrame(g3d_data)

    #save as g3d file
    result.to_csv(g3d_name_out, sep="\t", index=False)

if __name__ == "__main__":
    ARGS = get_cli_arguments()

    # Read Fasta file and get chromosome length
    CHROMOSOME_LENGTH = extract_chromosome_length(ARGS.fasta)

    # Assign chromosome number
    convert_to_g3d(ARGS.pdb, CHROMOSOME_LENGTH, ARGS.resolution, ARGS.output)

import argparse
from pathlib import Path

import jinja2


def is_file(parser, file_path):
    """Check file exists.

    Parameters
    ----------
    parser : argparse.ArgumentParser
        Command line argument parser.
    file_path : str
        File path
    Returns
    -------
    str
        File path.
    """
    if not Path(file_path).is_file():
        parser.error(f"The file {file_path} does not exist")
    else:
        return file_path


def get_cli_arguments():
    """Command line argument parser.

    Returns
    -------
    argparse.Namespace
        Object containing arguments.
    """
    parser = argparse.ArgumentParser(add_help=False)
    required = parser.add_argument_group("required arguments")
    optional = parser.add_argument_group("optional arguments")
    required.add_argument(
        "--template",
        action="store",
        type=lambda name: is_file(parser, name),
        help="Template file for HiC-Pro configuration.",
        required=True,
    )
    required.add_argument(
        "--chromosome-sizes",
        action="store",
        type=lambda name: is_file(parser, name),
        help="File with chromosome sizes.",
        required=True,
    )
    required.add_argument(
        "--genome-fragment",
        action="store",
        type=lambda name: is_file(parser, name),
        help="File with genome restriction sites.",
        required=True,
    )
    required.add_argument(
        "--genome-index-path",
        action="store",
        type=str,
        help="Path to bowtie2 indexes.",
        required=True,
    )
    required.add_argument(
        "--resolutions",
        action="store",
        type=int,
        nargs="+",
        help="HiC resolutions.",
        required=True,
    )
    required.add_argument(
        "--output",
        action="store",
        type=str,
        help="HiC-Pro configuration file.",
        required=True,
    )
    # Add help.
    optional.add_argument(
        "-h",
        "--help",
        action="help",
        default=argparse.SUPPRESS,
        help="Show this help message and exit.",
    )
    return parser.parse_args()


if __name__ == "__main__":
    ARGS = get_cli_arguments()

    with open(ARGS.template, "r") as template_file, \
        open(ARGS.output, "w") as out_file:
        template = jinja2.Template(template_file.read())
        out_file.write(
            template.render(
                genome_index_path=str(Path(ARGS.genome_index_path).resolve()) + "/",
                chromosome_sizes=Path(ARGS.chromosome_sizes).resolve(),
                genome_fragment=Path(ARGS.genome_fragment).resolve(),
                resolutions=" ".join([str(res) for res in ARGS.resolutions])
            )
        )

from watermark import watermark


def list_versions():
    """List package versions and describe environment."""
    print(
        watermark(
            machine=True,
            python=True,
            packages="numpy,pandas,scipy,biopandas,matplotlib,Bio,jinja2,sklearn,iced,pastis",
            watermark=True,
            conda=True
        )
    )


if __name__ == "__main__":
    list_versions()
    print("-" * 40)

import argparse
import numpy as np
import pandas as pd
from pathlib import Path

import iced
from iced import io
from pastis.optimization import mds
from pastis import dispersion
from pastis.optimization import negative_binomial_structure
from pastis.io.write_struct.pdb import writePDB
from scipy import sparse

def is_file(parser, file_path):
    """Check file exists.

    Parameters
    ----------
    parser : argparse.ArgumentParser
        Command line argument parser
    file_path : str
        File path
    Returns
    -------
    str
        File path    
    """
    if not Path(file_path).is_file():
        parser.error(f"The file {file_path} does not exist")
    else:
        return file_path


def get_cli_arguments():
    """Command line argument parser.

    Returns
    -------
    argparse.Namespace
        Object containing arguments
    """
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--matrix",
        action="store",
        type=lambda name: is_file(parser, name),
        help="Name of Matrix file",
        required=True,
    )
    parser.add_argument(
        "--bed",
        action="store",
        type=lambda name: is_file(parser, name),
        help="Name of Bed file",
        required=True,
    )
    parser.add_argument(
        "--output",
        action="store",
        type=str,
        help="Name of output PDB file containing the 3D structure of the genome",
        required=True,
    )
    return parser.parse_args()


def run_pastis_nb(matrix_filename, bed_filename, output_filename, seed=0, percentage_to_filter=0.02):
    """Build 3D structure of genome with Pastis NB algorithme.

    Parameters
    ----------
    matrix_filename : str
        Name of Matrix file
    bed_filename : str
        Name of Bed file
    output_filename : str
        Name of output PDB file with 3D structure of the genome
    seed : int (default: 0)
        Random generator seed
    percentage_to_filter : float (default: 0.02)
        Percentage to filter out data 
    """
    ###############################################################################
    #Load chromosome lengths and count data
    lengths = io.load_lengths(bed_filename)

    ###############################################################################
    # The sparse matrix generated by HiC-Pro seems to randomly miss a row.
    # The dense matrix generated by HiC-Pro_3.1.0/bin/utils/sparseToDense.py is always complete.
    # The 3D model is thus inferred from it.
    #counts = io.load_counts(matrix_filename, base=1)
    counts_maps = pd.read_csv(matrix_filename, sep='\t', header=None)
    counts_maps = counts_maps.values
    # to solve : File "iced/_filter_.pyx", line 15, in iced._filter_._filter_csr ValueError: Buffer dtype mismatch, expected 'DOUBLE' but got 'long'
    counts_maps = counts_maps.astype("double")
    #counts_maps = np.load(matrix_filename)
    counts_maps[np.isnan(counts_maps)] = 0
    counts = sparse.coo_matrix(counts_maps)
    counts.setdiag(0)
    counts.eliminate_zeros()
    counts = counts.tocoo()

    ###############################################################################
    # Normalize the data
    counts = iced.filter.filter_low_counts(
        counts, percentage=percentage_to_filter,
        sparsity=False)
    normed, bias = iced.normalization.ICE_normalization(
        counts, max_iter=300,
        output_bias=True)
    # Save the normalization data and biases just for other purposes
    # io.write_counts("./nb_normalized.matrix", normed)
    # np.savetxt("./nb.biases", bias)
    random_state = np.random.RandomState(seed)

    ###############################################################################
    # First estimate MDS for initialization
    X = mds.estimate_X(normed, random_state=random_state)
    ###############################################################################
    # Estimate constant dispersion parameters
    dispersion_ = dispersion.ExponentialDispersion(degree=0)
    _, mean, variance, weights = dispersion.compute_mean_variance(counts, lengths, bias=bias)
    dispersion_.fit(mean, variance, sample_weights=(weights**0.5))
    ###############################################################################
    # Now perform NB 3D inference.
    alpha = -3
    beta = 1
    counts = counts.tocoo()
    print("Estimating structure")
    X = negative_binomial_structure.estimate_X(
        counts, alpha, beta, bias=bias,
        lengths=lengths,
        dispersion=dispersion_,
        use_zero_entries=True,
        ini=X.flatten())
    ###############################################################################
    # Remove beads that were not infered
    mask = (np.array(counts.sum(axis=0)).flatten() + np.array(counts.sum(axis=1)).flatten() == 0)
    mask = mask.flatten()
    X_ = X.copy()
    X_[mask] = np.nan
    ###############################################################################
    np.savetxt(output_filename + ".txt", X_)
    print("Results written to", output_filename)
    writePDB(X_, output_filename)


if __name__ == "__main__":
    ARGS = get_cli_arguments()
    run_pastis_nb(
        matrix_filename=ARGS.matrix,
        bed_filename=ARGS.bed,
        output_filename=ARGS.output,
        seed=0, 
        percentage_to_filter=0.02
    )