#!/usr/bin/env python
"""
This is an executable file that accepts two conformations (old-style or URI), 
or a number of a block, or /path/to/the/trajectory/number   ( "mytraj/5" for example). 

Called in a folder with trajectory, "$dif 5 10" will calculate the rMSD between blocks 5 and 10. 

Three examples below would work for old style trajectory (1), both (2) and new style trajectory (3) 

$dif mytraj/block5.dat mytraj/block10.dat
$dif mytraj/5 mytraj/6 
$dif mytraj/blocks_5-10.h5::5 blocks_5010.h5::10 

It calculates sqrt(mean square displacement) after accounting for shifts and rotations of the 
conformation as a whole. 

It then prints a "shuffled distance", which is rMSD between the first conformation and the second conformation 
in which the monomer indices were shuffled. It is a measure of how much displacement is "possible" 
between the first and the second conformation. Specifically useful if you have 
spherical confinement - then shuffled distance becomes a proxy of how much would an MSD be 
between two random conformations inside a sphere. In that case, when rMSD approaches shuffled distance, 
you likely have mixed enough to reach some kind of equilibrium (and if you haven't you definitely haven't
reached equilibrium). 
"""
from __future__ import absolute_import, division, print_function

import os
import sys

import numpy as np
import scipy.optimize

from polychrom.polymer_analyses import kabsch_msd
from polychrom.polymerutils import fetch_block, load

filename = sys.argv[1]
try:
    fol, fname = os.path.split(filename)
    data1 = fetch_block(fol, fname)
except:
    data1 = load(filename)

filename = sys.argv[2]
try:
    fol, fname = os.path.split(filename)
    data2 = fetch_block(fol, fname)
except:
    data2 = load(filename)

dataa = data1 - np.sum(data1, 0) / float(len(data1))
datab = data2 - np.sum(data2, 0) / float(len(data2))


print(
    "Shuffled distance",
    np.mean(
        [
            np.sqrt(kabsch_msd(
                dataa, datab[np.argsort(np.random.random(len(datab)))]
            ))
            for _ in range(10)
        ]
    ),
)

print("actual distance", np.sqrt(kabsch_msd(dataa, datab)))
