import numpy as np
import pandas as pd
from matplotlib import pyplot as plt

from ovito.modifiers import SpatialCorrelationFunctionModifier
from ovito.io import import_file, export_file

dump_file = 'test.dump'

pipeline = import_file(dump_file)

mod = SpatialCorrelationFunctionModifier(
    direct_summation=True,
    neighbor_cutoff=46 / 2.0,
    property1='atomicvolume',
    property2='atomicvolume'
)
pipeline.modifiers.append(mod)
data = pipeline.compute()

Cr_direct = data.tables['correlation-neighbor'].xy()
rdf_direct = data.tables['correlation-neighbor-rdf'].xy()

C_norm = Cr_direct.copy()
mean1 = data.attributes['CorrelationFunction.mean1']
mean2 = data.attributes['CorrelationFunction.mean2']
covariance = data.attributes['CorrelationFunction.covariance']
C_norm[:,1] = (C_norm[:,1] - mean1*mean2) / (covariance - mean1*mean2)

plt.plot(C_norm[:, 0], C_norm[:, 1], label='Ovito-direct')

plt.legend(title=r'$C(r)$ normalize by covariance')
plt.yscale('log')
#plt.ylim(1e-2, 1e4)
plt.xlabel('r')
plt.ylabel('$C(r)$')