#!/usr/bin/env python3

import argparse as ap
import numpy as np
import pandas as pd
import pickle

from tqdm import tqdm

# Parse arguments
parser = ap.ArgumentParser()
parser.add_argument("files", nargs="+")
args = parser.parse_args()

# Open first file to get targets
print("Reading targets...", end="", flush=True)
targets = np.loadtxt(args.files[0], usecols=1, dtype="U4")
n_targets = len(targets)
print("done")

# Build DataFrame
# The DF are used to ensure that distances and ligand similarities are inserted at the correct place
# Initializing the DF with a numpy array is essential for speed at assignment
print("Allocating DataFrame memory...", end="", flush=True)
df_dist = pd.DataFrame(index=targets, columns=targets, data=-1*np.ones((n_targets,n_targets)))
df_lsim = pd.DataFrame(index=targets, columns=targets, data=-1*np.ones((n_targets,n_targets)))
print("done")

print("Merging data...", flush=True)
for fname in tqdm(args.files):
    target = np.loadtxt(fname, usecols=0, dtype="U4")[0]
    ctargets = np.loadtxt(fname, usecols=1, dtype="U4") # Can be removed if targets == ctargets all the time
    dist = np.loadtxt(fname, usecols=2)
    lsim = np.loadtxt(fname, usecols=3)
    
    # Populate distance matrix
    if len(dist) == n_targets:
        df_dist.loc[target, ctargets] = dist
    else:
        print("  Invalid number of distances for {target}")

    # Populate ligand similarity matrix
    if len(lsim) == n_targets:
        df_lsim.loc[target, ctargets] = lsim
    else:
        print("  Invalid number of ligand similarities for {target}")

dist = df_dist.values
lsim = df_lsim.values

# Check properties
print("Checking matrix properties...", flush=True)
assert np.allclose(np.diagonal(dist), np.zeros(n_targets))
assert np.allclose(np.diagonal(lsim), np.ones(n_targets))
assert np.allclose(dist, dist.T)
assert np.allclose(lsim, lsim)
print("done")

# Set NaNs for compatibility with original implementation
dist[dist < 0] = np.nan
lsim[lsim < 0] = np.nan

print("Checking data...", flush=True)
rows, cols = np.where(np.isnan(dist)) # Invalid distances
for t1, t2 in zip(df_dist.index.values[rows], df_dist.columns.values[cols]):
    print(f"  Missing distance for {t1} {t2}")
rows, cols = np.where(np.isnan(lsim)) # Invalid ligand similarities
for t1, t2 in zip(df_dist.index.values[rows], df_dist.columns.values[cols]):
    print(f"  Missing ligand similarity for {t1} {t2}")

print("Dumping pickle object...", end="", flush=True)
pickle.dump((dist, targets, lsim), open('matrix.pickle','wb'),-1)
print("done")
@@ -201,16 +201,76 @@
 args=%22+%22
+, type=str)%0Aparser.add_argument(%22-out%22, %22--output%22, type=str
 )%0Aargs =
@@ -2473,16 +2473,17 @@
 %0A%0Aprint(
+f
 %22Dumping
@@ -2496,16 +2496,27 @@
 e object
+ %7Bargs.out%7D
 ...%22, en
@@ -2577,28 +2577,27 @@
 pen(
-'matrix.pickle','wb'
+f%22%7Bargs.out%7D%22, %22wb%22
 ),-1

add output