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

This module contains a collection of functions for analyzing diffusion data.

Written by Vladislav Belyy (UCSF)

"""
import os
import numpy as np
import math
from scipy import stats
from scipy.optimize import curve_fit
import matplotlib.pylab as plt


### Definitions of short useful functions #####################################
def ax_fun(x, a):
    # Line with zero y-intercept
    return a*x

def get_residuals(data_x, data_y, slope, y_intercept):
    # Calculate residuals of a linear fit for x-y data
    model_y = data_x * slope + y_intercept
    resid = data_y - model_y
    return resid

def get_sum_squares(data):
    # Calculate total sum of squares
    mean = np.mean(data)
    sum_sq = np.sum(np.square(data - mean))
    return sum_sq

def get_r_sq(data_x, data_y, slope, yint):
    # Calculate r-squared statistic for a given line
    resid = get_residuals(data_x, data_y, slope, yint)
    sum_resid = np.sum(np.square(resid))
    sum_sq = get_sum_squares(data_y)
    r_sq = 1 - sum_resid/sum_sq
    return r_sq

### Definitions of major functions ############################################
    
def diffconst_gw(R, T, c, eta=0.056):
    """Calculate diffusion constant based on Guigas and Weiss 2006 paper.
    
    Returns the calculated diffusion constant of a membrane inclusion, as
    derived in the 2006 paper by Guigas and Weiss, 
    doi 10.1529/biophysj.106.087031 .
    
    Args:
        R (float): radius of inclusion, in nm.
        T (float): Temperature, in Kelvin.
        c (float): boundary conditions, in nm (c=6 for stick, c=4 for slip).
        eta (float): membrane viscosity, in Pa s. Defaults to 0.056 (the value
            used in the paper).
        
    Returns:
        D (float): Diffusion constant, in um^2/s.
 
    """
    
    k = 1.38e-23 # Boltzmann constant in J*K^(-1)
    R_m = R / 1e9 # convert nm to meters
    u_conv = 1e12 # to convert from m^2/s to um^2/s
    D = u_conv*k*T*math.atan(c/float(R)) / (8*np.pi *eta *R_m)
    return D

def diffconst_sd(R, T, c, eta_m=0.19, eta_c=0.039, h=3.5):
    """Calculate diffusion constant based on work by Saffman and Delbruck.