# -*- coding: utf-8 -*-
"""SASCalc

Automatically generated by Colaboratory.

Original file is located at
    https://colab.research.google.com/drive/16WemERFWTY_7Jb3FvuQf1bXT19B8-egU
"""

from __future__ import absolute_import, division, print_function, unicode_literals
from builtins import object, range, map, zip

# from io import open
"""
Created on December 12, 2015

@author: Jesse B. Hopkins

#******************************************************************************
# This file is part of RAW.
#
#    RAW is free software: you can redistribute it and/or modify
#    it under the terms of the GNU General Public License as published by
#    the Free Software Foundation, either version 3 of the License, or
#    (at your option) any later version.
#
#    RAW is distributed in the hope that it will be useful,
#    but WITHOUT ANY WARRANTY; without even the implied warranty of
#    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
#    GNU General Public License for more details.
#
#    You should have received a copy of the GNU General Public License
#    along with RAW.  If not, see <http://www.gnu.org/licenses/>.
#
#******************************************************************************

The purpose of this module is to contain functions for calculating
values from SAXS profiles. These are intended to be automated
functions, including calculation of rg and molecular weight.

It also contains functions for calling outside packages for use in RAW, like DAMMIF.
"""
try:
    import queue
except Exception:
    import Queue as queue

# import os
import time
import subprocess
import threading
import platform
import re
import math
import traceback
import copy
import tempfile

import numpy as np
from scipy import integrate as integrate
import scipy.interpolate
import scipy.signal
import scipy.stats as stats

# from scipy.constants import Avogadro
from numba import jit

# import sasexceptions as SASExceptions
from . import sasexceptions as SASExceptions

# import bioxtasraw.SASExceptions as SASExceptions
# import bioxtasraw.SASM as SASM
from . import sasm as SASM


# Define the rg fit function
@jit(nopython=True, cache=True, parallel=False)
def linear_func(x, a, b):
    return a + b * x


@jit(nopython=True, cache=True, parallel=False)
def weighted_lin_reg(x, y, err):
    weights = 1.0 / (err) ** 2.0

    w_sum = weights.sum()
    wy_sum = (weights * y).sum()
    wx_sum = (weights * x).sum()
    wxsq_sum = (weights * x**2.0).sum()
    wxy_sum = (weights * x * y).sum()

    delta = weights.sum() * wxsq_sum - (wx_sum) ** 2.0

    if delta != 0:
        a = (wxsq_sum * wy_sum - wx_sum * wxy_sum) / delta
        b = (w_sum * wxy_sum - wx_sum * wy_sum) / delta

        cov_a = wxsq_sum / delta
        cov_b = w_sum / delta
    else:
        a = -1
        b = -1
        cov_a = -1
        cov_b = -1

    return a, b, cov_a, cov_b


@jit(nopython=True, cache=True, parallel=False)
def lin_reg(x, y):
    x_sum = x.sum()
    xsq_sum = (x**2).sum()
    y_sum = y.sum()
    xy_sum = (x * y).sum()
    n = len(x)

    delta = n * xsq_sum - x_sum**2.0

    if delta != 0:
        a = (xsq_sum * y_sum - x_sum * xy_sum) / delta
        b = (n * xy_sum - x_sum * y_sum) / delta

        cov_y = (1.0 / (n - 2.0)) * ((y - a - b * x) ** 2.0).sum()
        cov_a = cov_y * (xsq_sum / delta)
        cov_b = cov_y * (n / delta)
    else:
        a = -1
        b = -1
        cov_a = -1
        cov_b = -1

    return a, b, cov_a, cov_b


@jit(nopython=True, cache=True, parallel=False)
def calcRg(q, i, err, transform=True, error_weight=True):
    if transform:
        # Start out by transforming as usual.
        x = np.square(q)
        y = np.log(i)
        yerr = np.absolute(
            err / i
        )  # I know it looks odd, but it's correct for a natural log
        x = x[np.where(np.isfinite(y))]
        yerr = yerr[np.where(np.isfinite(y))]
        y = y[np.where(np.isfinite(y))]
    else:
        x = q
        y = i
        yerr = err

    if error_weight:
        if np.any(yerr == 0):
            error_weight = False

    if error_weight:
        a, b, cov_a, cov_b = weighted_lin_reg(x, y, yerr)
    else:
        a, b, cov_a, cov_b = lin_reg(x, y)

    if b < 0:
        RG = np.sqrt(-3.0 * b)
        I0 = np.exp(a)

        # error in rg and i0 is calculated by noting that q(x)+/-Dq has Dq=abs(dq/dx)Dx, where q(x) is your function you're using
        # on the quantity x+/-Dx, with Dq and Dx as the uncertainties and dq/dx the derviative of q with respect to x.
        RGer = np.absolute(0.5 * (np.sqrt(-3.0 / b))) * np.sqrt(np.absolute(cov_b))
        I0er = I0 * np.sqrt(np.absolute(cov_a))

    else:
        RG = -1
        I0 = -1
        RGer = -1
        I0er = -1

    return RG, I0, RGer, I0er, a, b


def estimate_guinier_error(q, i, err, transform=True, error_weight=True):
    if transform:
        # Start out by transforming as usual.
        x = np.square(q)
        y = np.log(i)
        yerr = np.absolute(
            err / i
        )  # I know it looks odd, but it's correct for a natural log
    else:
        x = q
        y = i
        yerr = err

    win_size = len(x)

    if win_size < 10:
        est_rg_err = None
        est_i0_err = None
    else:
        var = win_size // 10
        if var > 12:
            step = int(np.ceil(var / 12.0))
        else:
            step = 1
        rg_list = []
        i0_list = []

        for li in range(0, var + 1, step):
            for ri in range(0, var + 1, step):
                if ri == 0:
                    Rg, I0, Rger, I0er, a, b = calcRg(
                        x[li:],
                        y[li:],
                        yerr[li:],
                        transform=transform,
                        error_weight=error_weight,
                    )
                else:
                    Rg, I0, Rger, I0er, a, b = calcRg(
                        x[li:-ri],
                        y[li:-ri],
                        yerr[li:-ri],
                        transform=transform,
                        error_weight=error_weight,
                    )

                rg_list.append(Rg)
                i0_list.append(I0)

        est_rg_err = np.array(rg_list).std()
        est_i0_err = np.array(i0_list).std()

    return est_rg_err, est_i0_err


##Porod Stuff####### uncommented 01/10/2023


def calcRefMW(i0, conc, ref_i0, ref_conc, ref_mw):
    if ref_mw > 0 and ref_i0 > 0 and ref_conc > 0 and conc > 0 and i0 > 0:
        mw = (i0 * (ref_mw / (ref_i0 / ref_conc))) / conc
    else:
        mw = -1

    return mw


def vpA(q_max):
    A = (
        -2.114 * 10**6 * q_max**4
        + 2.920 * 10**6 * q_max**3
        - 1.472 * 10**6 * q_max**2
        + 3.349 * 10**5 * q_max
        - 3.577 * 10**4
    )
    return A


def vpB(q_max):
    B = 12.09 * q_max**3 - 9.39 * q_max**2 + 3.03 * q_max + 0.29
    return B


def calcVqmax(q, i, rg, i0, choice="8/Rg", qmax=None):
    vpqmax = None

    if choice == "Default":
        if rg != 0:
            vpqmax = 8.0 / rg

            if vpqmax > 0.5 or vpqmax < 0.1:
                iratio = np.abs(np.log10(i0 / i) - 2.25)
                idx = np.argmin(iratio)

                vpqmax = q[idx]

            if vpqmax > 0.5:
                vpqmax = 0.5
            elif vpqmax < 0.1:
                vpqmax = 0.1

    elif choice == "8/Rg":
        if rg != 0:
            vpqmax = 8.0 / rg

            if vpqmax > 0.5:
                vpqmax = 0.5
            elif vpqmax < 0.1:
                vpqmax = 0.1

    elif choice == "log(I0/I(q))":
        if i0 != 0:
            iratio = np.abs(np.log10(i0 / i) - 2.25)
            idx = np.argmin(iratio)

            vpqmax = q[idx]

            if vpqmax > 0.5:
                vpqmax = 0.5
            elif vpqmax < 0.1:
                vpqmax = 0.1

    elif choice == "Manual":
        vpqmax = qmax

    if vpqmax is None:
        vpqmax = min(q[-1], 0.5)

    else:
        if vpqmax > q[-1]:
            vpqmax = q[-1]
        elif vpqmax < q[0]:
            vpqmax = q[0]
        else:
            idx = np.argmin(np.abs(q - vpqmax))
            vpqmax = q[idx]

    return vpqmax


def calcVpMW(q, i, err, rg, i0, rg_qmin, vp_density, qmax):
    # These functions are used to correct the porod volume for the length of the q vector

    if qmax not in q:
        idx = np.argmin(np.abs(q - qmax))
        qmax = q[idx]
    else:
        idx = np.argwhere(q == qmax)[0][0]

    q = q[: idx + 1]
    i = i[: idx + 1]
    err = err[: idx + 1]

    if q[-1] <= 0.5 and q[-1] >= 0.1:
        A = vpA(q[-1])
        B = vpB(q[-1])
    else:
        A = 0
        B = 1

    if i0 > 0:
        # Calculate the Porod Volume
        pVolume = porodVolume(q, i, err, rg, i0, interp=True, rg_qmin=rg_qmin)

        if pVolume == -1:
            mw = -1
            pv_cor = -1

        else:
            # Correct for the length of the q vector
            pv_cor = A + B * pVolume

            mw = pv_cor * vp_density

    else:
        mw = -1
        pVolume = -1
        pv_cor = -1

    return mw, pVolume, pv_cor


def calcAbsMW(i0, conc, rho_Mprot, rho_solv, nu_bar, r0):
    d_rho = (rho_Mprot - (rho_solv * nu_bar)) * r0
    mw = (Avogadro * i0 / conc) / np.square(d_rho)
    return mw


def volumeOfCorrelation(q, i, i0):
    """Calculates the volume of correlation as the ratio of i0 to $\int q*I dq$"""
    tot = integrate.trapz(q * i, q)
    vc = i0 / tot
    return vc


def porodInvariant(q, i, start=0, stop=-1):
    return integrate.trapz(i[start:stop] * np.square(q[start:stop]), q[start:stop])


def porodVolume(q, i, err, rg, i0, start=0, stop=-1, interp=True, rg_qmin=0):
    if interp and q[0] != 0:

        def f(x):
            return i0 * np.exp((-1.0 / 3.0) * np.square(rg) * np.square(x))

        if rg_qmin > 0:
            findClosest = lambda a, l: min(l, key=lambda x: abs(x - a))
            closest_qmin = findClosest(rg_qmin, q)

            idx_min = np.where(q == closest_qmin)[0][0]

            q = q[idx_min:]
            i = i[idx_min:]
            err = err[idx_min:]

        if len(q) != 1:
            q_interp = np.arange(0, q[0], q[1] - q[0])
            i_interp = f(q_interp)
            err_interp = np.sqrt(i_interp)

            q = np.concatenate((q_interp, q))
            i = np.concatenate((i_interp, i))
            err = np.concatenate((err_interp, err))

    if len(q) != 1:
        pInvar = porodInvariant(q, i, start, stop)

        pVolume = 2 * np.pi**2 * i0 / pInvar

    else:
        pVolume = -1

    return pVolume


# Uncomented 01/10/2023 (bottom)


def autoRg(sasm, single_fit=False, error_weight=True):
    # This function automatically calculates the radius of gyration and scattering intensity at zero angle
    # from a given scattering profile. It roughly follows the method used by the autorg function in the atsas package

    q = sasm.getQ()
    i = sasm.getI()
    err = sasm.getErr()

    qmin = 0
    # 1
    try:
        rg, rger, i0, i0er, idx_min, idx_max = autoRg_inner(
            q,
            i,
            err,
            qmin,
            single_fit,
            error_weight,
            min_window=10,
            min_qrg=1.0,
            max_qrg=1.35,
            quality_thresh=0.6,
            data_range_scale=0,
            corr_coefht=2.0,
            win_length_weight=1.0,
        )
    except Exception:  # Catches unexpected numba errors, I hope
        # ** Interesting, guess just if anything in the exceptions class is called (maybe lin fit can give wrong math easy
        # e.g division by zero? Revisit to understand why this works and how well does it)
        traceback.print_exc()
        rg = -1
        rger = -1
        i0 = -1
        i0er = -1
        idx_min = -1
        idx_max = -1

    if rg == -1:
        # If we don't find a fit, relax the criteria
        try:
            rg, rger, i0, i0er, idx_min, idx_max = autoRg_inner(
                q,
                i,
                err,
                qmin,
                single_fit,
                error_weight,
                min_window=5,
                min_qrg=1.0,
                max_qrg=1.35,
                quality_thresh=0.5,
                data_range_scale=0,
                corr_coefht=2.0,
                win_length_weight=1.0,
            )
        except Exception:  # Catches unexpected numba errors, I hope
            traceback.print_exc()
            rg = -1
            rger = -1
            i0 = -1
            i0er = -1
            idx_min = -1
            idx_max = -1

    if rg == -1:
        # If we don't find a fit, relax the criteria
        try:
            rg, rger, i0, i0er, idx_min, idx_max = autoRg_inner(
                q,
                i,
                err,
                qmin,
                single_fit,
                error_weight,
                min_window=10,
                min_qrg=1.0,
                max_qrg=1.35,
                quality_thresh=0.6,
                data_range_scale=100,
                corr_coefht=2.0,
                win_length_weight=1.0,
            )
        except Exception:  # Catches unexpected numba errors, I hope
            traceback.print_exc()
            rg = -1
            rger = -1
            i0 = -1
            i0er = -1
            idx_min = -1
            idx_max = -1

    if rg == -1:
        # If we don't find a fit, relax the criteria
        try:
            rg, rger, i0, i0er, idx_min, idx_max = autoRg_inner(
                q,
                i,
                err,
                qmin,
                single_fit,
                error_weight,
                min_window=10,
                min_qrg=1.2,
                max_qrg=1.5,
                quality_thresh=0.3,
                data_range_scale=100,
                corr_coefht=2.0,
                win_length_weight=1.0,
            )
        except Exception:  # Catches unexpected numba errors, I hope
            traceback.print_exc()
            rg = -1
            rger = -1
            i0 = -1
            i0er = -1
            idx_min = -1
            idx_max = -1

    if rg == -1:
        # If we don't find a fit, relax the criteria
        try:
            rg, rger, i0, i0er, idx_min, idx_max = autoRg_inner(
                q,
                i,
                err,
                qmin,
                single_fit,
                error_weight,
                min_window=5,
                min_qrg=1.2,
                max_qrg=1.5,
                quality_thresh=0.3,
                data_range_scale=100,
                corr_coefht=2.0,
                win_length_weight=1.0,
            )
        except Exception:  # Catches unexpected numba errors, I hope
            traceback.print_exc()
            rg = -1
            rger = -1
            i0 = -1
            i0er = -1
            idx_min = -1
            idx_max = -1

    return rg, rger, i0, i0er, idx_min, idx_max


@jit(nopython=True, cache=True, parallel=False)
def autoRg_inner(
    q,
    i,
    err,
    qmin,
    single_fit,
    error_weight,
    min_window=10,
    min_qrg=1.0,
    max_qrg=1.35,
    quality_thresh=0.6,
    data_range_scale=0,
    corr_coefht=2.0,
    win_length_weight=1.0,
):
    # Pick the start of the RG fitting range. Note that in autorg, this is done
    # by looking for strong deviations at low q from aggregation or structure factor
    # or instrumental scattering, and ignoring those. This function isn't that advanced
    # so we start at 0.

    # Note, in order to speed this up using numba, I had to do some unpythonic things
    # with declaring lists ahead of time, and making sure lists didn't have multiple
    # object types in them. It makes the code a bit more messy than the original
    # version, but numba provides a significant speedup.

    # Have to pick the right Starting range to avoid various weirdnesses in the data
    data_start = (i > 0).argmax()  # **Remember first I value is highest usually

    if len(i) > 20:
        while i[data_start] < np.mean(i[-1 * int(len(i) / 20) :]) and data_start < len(
            i
        ):
            data_start = (i[data_start + 1 :] > 0).argmax() + data_start + 1

        while i[data_start] > np.mean(i[data_start : data_start + 20]) * 10:
            data_start = (i[data_start + 1 :] > 0).argmax() + data_start + 1

    # Turns out to be pretty important to pick a good initial q range for the search
    # This is just kind of determined by what looks reasonable
    if data_range_scale == 0:
        total_int_range = np.abs(
            np.mean(i[data_start : data_start + 20]) / np.mean(i[-20:])
        )
        if total_int_range < 20:
            data_range_scale = 2.5

        elif total_int_range < 100:
            data_range_scale = 5

        elif total_int_range < 1000:
            data_range_scale = 10

        else:
            data_range_scale = 100

    ##Following the atsas package, the end point of our search space is the q value
    # where the intensity has droped by an order of magnitude from the initial value.
    data_end = np.abs(i[i > 0] - i[data_start] / data_range_scale).argmin()

    i_val = i[i > 0][data_end]
    data_end = np.argwhere(i == i_val)[0][0]

    # This makes sure we're not getting some weird fluke at the end of the scattering profile.
    if data_end > len(i) / 2.0:
        found = False
        if len(i) > data_start + 20 and data_start + 20 < len(i) / 2.0:
            idx = data_start + 20
        else:
            idx = data_start
        while not found:
            idx = idx + 1
            if i[idx] < i[data_start] / data_range_scale:
                found = True
            elif idx == len(q) - 1:
                found = True
        data_end = idx

    # Start out by transforming as usual.
    qs = np.square(q)
    il = np.log(i)
    iler = np.absolute(err / i)

    # Pick a minimum fitting window size. 10 is consistent with atsas autorg.
    min_window = min_window

    max_window = data_end - data_start

    if max_window < min_window:
        max_window = min_window

    # It is very time consuming to search every possible window size and every possible starting point.
    # Here we define a subset to search.
    tot_points = max_window
    window_step = min(tot_points // 10, 20)
    data_step = tot_points // 50

    if window_step == 0:
        window_step = 1
    if data_step == 0:
        data_step = 1

    window_list = [
        0
        for k in range(
            int(math.ceil((max_window - min_window) / float(window_step))) + 1
        )
    ]

    for k in range(int(math.ceil((max_window - min_window) / float(window_step)))):
        window_list[k] = min_window + k * window_step

    window_list[-1] = max_window

    num_fits = 0

    for w in window_list:
        num_fits = num_fits + int(
            math.ceil((data_end - w - data_start) / float(data_step))
        )

    if num_fits < 0:
        num_fits = 1

    start_list = [0 for k in range(num_fits)]
    w_list = [0 for k in range(num_fits)]
    q_start_list = [0.0 for k in range(num_fits)]
    q_end_list = [0.0 for k in range(num_fits)]
    rg_list = [0.0 for k in range(num_fits)]
    rger_list = [0.0 for k in range(num_fits)]
    i0_list = [0.0 for k in range(num_fits)]
    i0er_list = [0.0 for k in range(num_fits)]
    qrg_start_list = [0.0 for k in range(num_fits)]
    qrg_end_list = [0.0 for k in range(num_fits)]
    rsqr_list = [0.0 for k in range(num_fits)]
    chi_sqr_list = [0.0 for k in range(num_fits)]
    reduced_chi_sqr_list = [0.0 for k in range(num_fits)]
    corr_coef_list = [0.0 for k in range(num_fits)]

    success = np.zeros(num_fits)

    current_fit = 0
    # This function takes every window size in the window list, stepts it through the data range, and
    # fits it to get the RG and I0. If basic conditions are met, qmin*RG<1 and qmax*RG<1.35, and RG>0.1,
    # We keep the fit.
    for w in window_list:
        for start in range(data_start, data_end - w, data_step):
            x = qs[start : start + w]
            y = il[start : start + w]
            yerr = iler[start : start + w]

            # Remove NaN and Inf values:
            x = x[np.where(np.isfinite(y))]
            yerr = yerr[np.where(np.isfinite(y))]
            y = y[np.where(np.isfinite(y))]

            RG, I0, RGer, I0er, a, b = calcRg(
                x, y, yerr, transform=False, error_weight=error_weight
            )

            if (
                RG > 0.1
                and q[start] * RG < min_qrg
                and q[start + w - 1] * RG < max_qrg
                and RGer / RG <= 1
            ):
                residual = il[start : start + w] - linear_func(
                    qs[start : start + w], a, b
                )

                r_sqr = (
                    1
                    - np.square(residual).sum()
                    / np.square(
                        il[start : start + w] - il[start : start + w].mean()
                    ).sum()
                )

                if r_sqr > 0.15:
                    chi_sqr = np.square((residual) / iler[start : start + w]).sum()

                    # All of my reduced chi_squared values are too small, so I suspect something isn't right with that.
                    # Values less than one tend to indicate either a wrong degree of freedom, or a serious overestimate
                    # of the error bars for the system.
                    dof = w - 2.0
                    reduced_chi_sqr = chi_sqr / dof

                    # Ideally this would be a pvalue, but I'd have to invest in a lot of intrastructure to actually calculate that in a jitted function
                    corr_coef = 1 - spearmanr(residual, qs[start : start + w])

                    start_list[current_fit] = start
                    w_list[current_fit] = w
                    q_start_list[current_fit] = q[start]
                    q_end_list[current_fit] = q[start + w - 1]
                    rg_list[current_fit] = RG
                    rger_list[current_fit] = RGer
                    i0_list[current_fit] = I0
                    i0er_list[current_fit] = I0er
                    qrg_start_list[current_fit] = q[start] * RG
                    qrg_end_list[current_fit] = q[start + w - 1] * RG
                    rsqr_list[current_fit] = r_sqr
                    chi_sqr_list[current_fit] = chi_sqr
                    reduced_chi_sqr_list[current_fit] = reduced_chi_sqr
                    corr_coef_list[current_fit] = corr_coef

                    success[current_fit] = 1

            current_fit = current_fit + 1

    if np.sum(success) > 0:
        fit_array = np.array(
            [
                [
                    start_list[k],
                    w_list[k],
                    q_start_list[k],
                    q_end_list[k],
                    rg_list[k],
                    rger_list[k],
                    i0_list[k],
                    i0er_list[k],
                    qrg_start_list[k],
                    qrg_end_list[k],
                    rsqr_list[k],
                    chi_sqr_list[k],
                    reduced_chi_sqr_list[k],
                    corr_coef_list[k],
                ]
                for k in range(num_fits)
                if success[k] == 1
            ]
        )

        # Now we evaluate the quality of the fits based both on fitting data and on other criteria.

        # Choice of weights is pretty arbitrary, but has been tested against
        # all the data in the SASBDB (as of 11/2020)
        qmaxrg_weight = 1
        qminrg_weight = 1
        rg_frac_err_weight = 1
        i0_frac_err_weight = 1
        r_sqr_weight = 4
        reduced_chi_sqr_weight = 0
        window_size_weight = win_length_weight
        corr_coef_weight = corr_coefht

        weights = np.array(
            [
                qmaxrg_weight,
                qminrg_weight,
                rg_frac_err_weight,
                i0_frac_err_weight,
                r_sqr_weight,
                reduced_chi_sqr_weight,
                window_size_weight,
                corr_coef_weight,
            ]
        )

        quality = np.zeros(len(fit_array))

        max_window_real = float(max(w_list))

        # This iterates through all the fits, and calculates a score. The score is out of 1, 1 being the best, 0 being the worst.
        indices = list(range(len(fit_array)))
        for a in indices:
            k = int(a)  # This is stupid and should not be necessary. Numba bug?

            # Scores all should be 1 based. Reduced chi_square score is not, hence it not being weighted.
            qmaxrg_score = 1 - abs((fit_array[k, 9] - 1.3) / 1.3)
            qminrg_score = 1 - fit_array[k, 8]
            rg_frac_err_score = 1 - fit_array[k, 5] / fit_array[k, 4]
            i0_frac_err_score = 1 - fit_array[k, 7] / fit_array[k, 6]
            r_sqr_score = fit_array[k, 10]
            reduced_chi_sqr_score = 1 / fit_array[k, 12]  # Not right
            window_size_score = fit_array[k, 1] / max_window_real
            corr_coef_score = fit_array[k, 13]

            scores = np.array(
                [
                    qmaxrg_score,
                    qminrg_score,
                    rg_frac_err_score,
                    i0_frac_err_score,
                    r_sqr_score,
                    reduced_chi_sqr_score,
                    window_size_score,
                    corr_coef_score,
                ]
            )

            total_score = (weights * scores).sum() / weights.sum()

            quality[k] = total_score

            # all_scores[k] = scores

        # I have picked an aribtrary threshold here. Not sure if 0.6 is a good quality cutoff or not.
        if quality.max() > quality_thresh:
            if not single_fit:
                idx = quality.argmax()
                rger = fit_array[:, 5][quality > quality[idx] - 0.1].std()
                i0er = fit_array[:, 7][quality > quality[idx] - 0.1].std()
                idx_min = int(fit_array[idx, 0])
                idx_max = int(fit_array[idx, 0] + fit_array[idx, 1] - 1)
            else:
                idx = quality.argmax()
                idx_min = int(fit_array[idx, 0])
                idx_max = int(fit_array[idx, 0] + fit_array[idx, 1] - 1)

            # Now refine the range a bit
            max_quality = quality.max()
            qual = max_quality

            idx_max_ref = idx_max

            if max_qrg == 1.35:
                max_qrg_ref = 1.3
            else:
                max_qrg_ref = max_qrg

            if q[idx_max] * RG < 1.0:
                quality_scale = 0.9
                r_thresh = 0.1
            else:
                quality_scale = 0.97
                r_thresh = 0.15

            # Refine upper end of range
            while qual > quality_scale * max_quality and idx_max_ref < len(q):
                idx_max_ref = idx_max_ref + 1
                x = qs[idx_min : idx_max_ref + 1]
                y = il[idx_min : idx_max_ref + 1]
                yerr = iler[idx_min : idx_max_ref + 1]

                # Remove NaN and Inf values:
                x = x[np.where(np.isfinite(y))]
                yerr = yerr[np.where(np.isfinite(y))]
                y = y[np.where(np.isfinite(y))]

                RG, I0, RGer, I0er, a, b = calcRg(
                    x, y, yerr, transform=False, error_weight=error_weight
                )

                if (
                    RG > 0.1
                    and q[idx_min] * RG < min_qrg
                    and q[idx_max_ref] * RG < max_qrg_ref
                    and RGer / RG <= 1
                ):
                    residual = il[idx_min : idx_max_ref + 1] - linear_func(
                        qs[idx_min : idx_max_ref + 1], a, b
                    )

                    r_sqr = (
                        1
                        - np.square(residual).sum()
                        / np.square(
                            il[idx_min : idx_max_ref + 1]
                            - il[idx_min : idx_max_ref + 1].mean()
                        ).sum()
                    )

                    if r_sqr > r_thresh:
                        chi_sqr = np.square(
                            residual / iler[idx_min : idx_max_ref + 1]
                        ).sum()

                        # All of my reduced chi_squared values are too small, so I suspect something isn't right with that.
                        # Values less than one tend to indicate either a wrong degree of freedom, or a serious overestimate
                        # of the error bars for the system.
                        dof = w - 2.0
                        reduced_chi_sqr = chi_sqr / dof

                        corr_coef = 1 - spearmanr(
                            residual, qs[idx_min : idx_max_ref + 1]
                        )

                        qmaxrg_score = 1 - abs((q[idx_max_ref] * RG - 1.3) / 1.3)
                        qminrg_score = 1 - q[idx_min] * RG
                        rg_frac_err_score = 1 - RGer / RG
                        i0_frac_err_score = 1 - I0er / I0
                        r_sqr_score = r_sqr
                        reduced_chi_sqr_score = 1 / reduced_chi_sqr  # Not right
                        window_size_score = fit_array[k, 1] / max_window_real
                        corr_coef_score = corr_coef

                        scores = np.array(
                            [
                                qmaxrg_score,
                                qminrg_score,
                                rg_frac_err_score,
                                i0_frac_err_score,
                                r_sqr_score,
                                reduced_chi_sqr_score,
                                window_size_score,
                                corr_coef_score,
                            ]
                        )

                        qual = (weights * scores).sum() / weights.sum()

                        if q[idx_max] * RG < 1.0:
                            quality_scale = 0.9
                            r_thresh = 0.1
                        else:
                            quality_scale = 0.97
                            r_thresh = 0.15

                    else:
                        qual = -1

                else:
                    qual = -1

                max_quality = max(max_quality, qual)

            idx_max = idx_max_ref - 1

            # Refine lower end of range
            idx_min_ref = idx_min
            qual = max_quality

            while qual > 0.97 * max_quality and idx_min_ref > 0:
                idx_min_ref = idx_min_ref - 1
                x = qs[idx_min_ref : idx_max + 1]
                y = il[idx_min_ref : idx_max + 1]
                yerr = iler[idx_min_ref : idx_max + 1]

                # Remove NaN and Inf values:
                x = x[np.where(np.isfinite(y))]
                yerr = yerr[np.where(np.isfinite(y))]
                y = y[np.where(np.isfinite(y))]

                RG, I0, RGer, I0er, a, b = calcRg(
                    x, y, yerr, transform=False, error_weight=error_weight
                )

                if (
                    RG > 0.1
                    and q[idx_min_ref] * RG < min_qrg
                    and q[idx_max] * RG < max_qrg_ref
                    and RGer / RG <= 1
                ):
                    residual = il[idx_min_ref : idx_max + 1] - linear_func(
                        qs[idx_min_ref : idx_max + 1], a, b
                    )

                    r_sqr = (
                        1
                        - np.square(residual).sum()
                        / np.square(
                            il[idx_min_ref : idx_max + 1]
                            - il[idx_min_ref : idx_max + 1].mean()
                        ).sum()
                    )

                    if r_sqr > 0.15:
                        chi_sqr = np.square(
                            (residual) / iler[idx_min_ref : idx_max + 1]
                        ).sum()

                        # All of my reduced chi_squared values are too small, so I suspect something isn't right with that.
                        # Values less than one tend to indicate either a wrong degree of freedom, or a serious overestimate
                        # of the error bars for the system.
                        dof = w - 2.0
                        reduced_chi_sqr = chi_sqr / dof

                        corr_coef = 1 - spearmanr(
                            residual, qs[idx_min_ref : idx_max + 1]
                        )

                        qmaxrg_score = 1 - abs((q[idx_max] * RG - 1.3) / 1.3)
                        qminrg_score = 1 - q[idx_min_ref] * RG
                        rg_frac_err_score = 1 - RGer / RG
                        i0_frac_err_score = 1 - I0er / I0
                        r_sqr_score = r_sqr
                        reduced_chi_sqr_score = 1 / reduced_chi_sqr  # Not right
                        window_size_score = fit_array[k, 1] / max_window_real
                        corr_coef_score = corr_coef

                        scores = np.array(
                            [
                                qmaxrg_score,
                                qminrg_score,
                                rg_frac_err_score,
                                i0_frac_err_score,
                                r_sqr_score,
                                reduced_chi_sqr_score,
                                window_size_score,
                                corr_coef_score,
                            ]
                        )

                        qual = (weights * scores).sum() / weights.sum()
                    else:
                        qual = -1

                else:
                    qual = -1

                max_quality = max(max_quality, qual)

            if idx_min_ref == 0 and qual != -1:
                idx_min = idx_min_ref
            else:
                idx_min = idx_min_ref + 1

            # Recalculate Guinier values with the new min and max idx
            x = qs[idx_min : idx_max + 1]
            y = il[idx_min : idx_max + 1]
            yerr = iler[idx_min : idx_max + 1]

            # Remove NaN and Inf values:
            x = x[np.where(np.isfinite(y))]
            yerr = yerr[np.where(np.isfinite(y))]
            y = y[np.where(np.isfinite(y))]

            rg, i0, Rger, I0er, a, b = calcRg(
                x, y, yerr, transform=False, error_weight=error_weight
            )

            if single_fit:
                rger = Rger
                i0er = I0er

        else:
            rg = -1
            rger = -1
            i0 = -1
            i0er = -1
            idx_min = -1
            idx_max = -1

    else:
        rg = -1
        rger = -1
        i0 = -1
        i0er = -1
        idx_min = -1
        idx_max = -1

    idx_min = idx_min + qmin
    idx_max = idx_max + qmin

    # returns Rg, Rg error, I0, I0 error, the index of the first q point of the fit and the index of the last q point of the fit
    return rg, rger, i0, i0er, idx_min, idx_max


@jit(nopython=True, cache=True)
def rankdata(array):
    temp = array.argsort()
    ranks = np.empty_like(temp)
    ranks[temp] = np.arange(len(array))

    return ranks


@jit(nopython=True, cache=True)
def spearmanr(array1, array2):
    rank1 = rankdata(array1)
    rank2 = rankdata(array2)

    n = rank1.size

    dsq = (rank1 - rank2) ** 2

    rho = 1.0 - (6 * dsq.sum()) / (n * (n**2 - 1))

    # t = rho*np.sqrt((n-2)/(1-rho**2)) #Calcuates t value for student t test

    return rho
