# -*- coding: utf-8 -*-
"""
Created on Mon Feb 19 15:33:05 2024

@author: regin
"""

import numpy as np
import matplotlib.pyplot as plt

gamma=7.
Omega=1
mu=1.
a=1.0
kbt=1
D=mu*kbt
N=2.
#t_prime=Drot/Omega**2
#x_prime= Drot/Omega
x_min=a*np.arcsin(Omega/gamma)
def force(x,Omega,gamma):
    return (Omega-gamma*np.sin(x/a))/mu
def potential(x,Omega,gamma):
    return -(Omega*x+gamma*a*np.cos(x/a))/mu
x=np.arange(-np.pi,np.pi*2*a*N,0.1)
x_harm=np.arange(-np.pi/1.5,np.pi/1.5,0.1)

from matplotlib.ticker import FuncFormatter, MultipleLocator



from mpl_toolkits.axes_grid1.inset_locator import inset_axes

def plot_washboard_with_HP():
    N=1

    Omega=1
    mu=1.
    a=1.0
    gamma=2.0
    x_min=a*np.arcsin(Omega/gamma)
    x=np.arange(-3/2*np.pi*a,2*np.pi*a,0.1)
    x_harm=np.arange(-np.pi/1.5,np.pi/1.5,0.1)
    x_harm2=np.arange(-0.6,0,0.1)
    x_harm3=np.arange(2*np.pi-0.6,2*np.pi,0.1)
    fig,ax = plt.subplots(1, 1, figsize=(3.375, 2.5)) 
    ax.plot(x/a,potential(x,Omega,0.5),label=r"$0.5$",linestyle="-",color='purple')
    ax.plot(x/a,potential(x,Omega,1),label=r"$1.0$",linestyle="-",color='red')
    ax.plot(x/a,potential(x,Omega,gamma),label=r"$2.0$",color='blue')
    ax.set_xlabel(r'$\vartheta $')
    ax.set_ylabel(r"$ U(\vartheta)D_\mathrm{rot}/k_\mathrm{B}T \omega$")

    ax.plot(x_harm,(x_harm-x_min)**2/(2*a)*gamma*np.cos(x_min/a)+potential(x_min,Omega,gamma),linestyle="--",color='black',label=r"HA")
    diff=potential(2*np.pi,Omega,gamma)-potential(0,Omega,gamma)
    #periodic
    #ax.plot(x_harm3,diff+(x_harm2-x_min)**2/(2*a)*gamma*np.cos(x_min/a)+potential(x_min,Omega,gamma),linestyle="--",color='black')

    ax.set_xticks((-5*np.pi/4, -3*np.pi/4, -np.pi/4 ,np.pi/4, 3*np.pi/4, 5*np.pi/4, 7*np.pi/4),minor=True)
    labels = [r'$-3\pi/2$', r'$-\pi$',r'$-\pi/2$' ,'$0$', r'$\pi/2$' ,r'$\pi$',  r'$3\pi/2$', r'$2\pi$']
    ax.set_xticks((-3*np.pi/2,-np.pi,-np.pi/2,0,np.pi/2,np.pi,3*np.pi/2,2*np.pi),labels,minor=False)
    ax.tick_params(which='both', bottom=True, top=True, left=True, right=True)
    
    # ax.set_xticklabels(labels)
    #ax.xaxis.set_major_formatter(plt.FuncFormatter(multiple_formatter()))
    ax.legend(title=r"$\gamma/\omega$")
    
        
    x=np.arange(0,np.pi,0.1)    
    x_harm=np.arange(0,np.pi/1.5,0.1)
    # Create inset
    inset_ax = inset_axes(ax, width="45%", height="35%", loc='lower left')
    inset_ax.plot(x / a, potential(x, Omega, 0.5), linestyle="-", color='purple')
    inset_ax.plot(x / a, potential(x, Omega, 1), linestyle="-", color='red')
    inset_ax.plot(x / a, potential(x, Omega, gamma), color='blue')
    inset_ax.plot(x_harm, (x_harm - x_min)**2 / (2 * a) * gamma * np.cos(x_min / a) + potential(x_min, Omega, gamma),
                  linestyle="--", color='black')
    
    # Customize inset ticks and labels
    inset_ax.tick_params(which='both', bottom=True, top=True, left=True, right=True)
    #inset_ax.set_xticks((np.pi/4, 3*np.pi/4, 5*np.pi/4, 7*np.pi/4),minor=True)
    #labels = ['$0$', r'$\pi/2$' ,r'$\pi$',  r'$3\pi/2$',r'$2\pi$']
    #inset_ax.set_xticks(( 0,np.pi/2,np.pi,3*np.pi/2, 2*np.pi),labels,minor=False)
    inset_ax.set_xticks([])
    inset_ax.set_yticks([])
        
    
    plt.tight_layout()
    #plt.savefig('../figures/washboard_potential.pdf')
    #plt.savefig('../figures/washboard_potential_inset.pdf')
    plt.show()

plot_washboard_with_HP()


def peq(x,Omega,gamma):
    return np.exp((Omega*x+gamma*a*np.cos(x/a)/D))
x_peq=np.arange(0,2*np.pi*N*a,0.1)


# Our integral approximation function
def integral_approximation(f, a, b):
    return (b-a)*np.mean(f)


# Define bounds of integral
an = 0.0
bn = 2.0*np.pi*a*N
# Generate function values
x_range = np.arange(an,bn+0.00001,.00001)
fx = peq(x_range, Omega, gamma)
# Approximate integral
Z = integral_approximation(fx,an,bn)

print(Z)
#plt.plot(x_peq/a,peq(x_peq,Omega, gamma)/(Z),label=r"$0.5$")
#plt.plot(x/a,potential(x,Omega, gamma),label=r"$1.5$")

#-------------------




def presentation_potential():
        
    def potential4(x,K):
        return K*x/x;
    
    def xn(n):
        return np.linspace((n-1)*np.pi,(n)*np.pi,100)
    
    def stepn(n):
        return [n*np.pi,n*np.pi+0.001]
    
    #plt.figure(figsize=(6, 2))
    
    plt.plot(xn(0),potential4(xn(0),0.5),label=r'$U_2$', color='red')
    plt.plot(xn(1),potential4(xn(1),-0.5),label=r'$U_2$', color='red')
    plt.plot(xn(2),potential4(xn(2),0.5),label=r'$U_2$', color='red')
    
    plt.plot(stepn(-1),[-0.5,0.5],label=r'$U_2$', color='red')
    plt.plot(stepn(0),[-0.5,0.5],label=r'$U_2$', color='red')
    plt.plot(stepn(1),[-0.5,0.5],label=r'$U_2$', color='red')
    plt.plot(stepn(2),[-0.5,0.5],label=r'$U_2$', color='red')
    plt.ylabel(r'$U(x)/k_bT$')
    #plt.xlim(-L/2,L/2)
    #plt.ylim(0,5)
    
    plt.xlabel(r'$x$')
    plt.xticks( [ -np.pi, 0, np.pi,2*np.pi],
                [r'$-a/2$',r'0',r'$a/2$', r'$a$']
        )
    # Auch Ticks für die y-Achse anpassen:
    plt.yticks( [-0.5, 0, 0.5],
                [ r'$U_1$', r'', r'$U_2$']
        )
    plt.tight_layout()
    #plt.savefig('C:/Users/Regin/OneDrive/Paper_tilted_washboard/presentations/presentation_01_23/potential_step.pdf')
    plt.show()
    
    x_gravitaxis=np.arange(0,2*np.pi,0.1)
    plt.plot(x_gravitaxis,potential(x_gravitaxis,1,1), color='green')
    plt.xlabel(r'$\vartheta$')
    plt.xticks( [0, np.pi,2*np.pi],
                [r'0',r'$a/2$', r'$a$']
        )
    plt.ylabel(r'$U(\vartheta)/k_bT$')
    plt.tight_layout()
    #plt.savefig('C:/Users/Regin/OneDrive/Paper_tilted_washboard/presentations/presentation_01_23/potential_gravitaxis.pdf')
    plt.show()
    
    
    
    x_gravitaxis=np.arange(-2*np.pi,2*np.pi,0.1)
    plt.plot(x_gravitaxis,potential(x_gravitaxis,1,1), color='blue')
    plt.xlabel(r'$x$')
    plt.xticks( [-2*np.pi,-np.pi,0,np.pi,2*np.pi],
                [r'$-a$',r'$-a/2$',r'0',r'$a/2$', r'$a$']
        )
    plt.ylabel(r'$U(x)/k_bT$')
    plt.tight_layout()
    #plt.savefig('C:/Users/Regin/OneDrive/Paper_tilted_washboard/presentations/presentation_01_23/potential_washboard.pdf')
    plt.show()
    return