# -*- coding: utf-8 -*-
"""
Created on Thu Jan 15 16:00:40 2026

@author: regin
"""

import numpy as np
import matplotlib.pyplot as plt
from nice_figures import *
figsize = load_style()
cols = load_cols()
colors=[cols['blue'],cols['orange'],cols['green'],cols['red'],cols['purple'],cols['brown'],cols['pink'],cols['olive']]
w, h = figsize[('APS', '1-column')]


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,a):
    return -(Omega*x+gamma*np.cos(x))
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

colors=['r', '#009900', 'b', '#bf00ff','#ffc40c','#e6550d']
from mpl_toolkits.axes_grid1.inset_locator import inset_axes

def plot_washboard_with_HP():
    N=1
    shift=np.pi+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,2*np.pi,0.1)
    x_harm=np.arange(-np.pi/1.5,np.pi/1.5,0.1)
    
    fig, ax = plt.subplots(figsize=(1.2 * w, h))
    ax.plot(x/a,potential(x,Omega,0.5,a)-shift,label=r"$2.0$",linestyle="-",color=colors[5])
    ax.plot(x/a,potential(x,Omega,1,a)-shift,label=r"$1.0$",linestyle="-",color=colors[3])
    ax.plot(x/a,potential(x,Omega,gamma,a)-shift,label=r"$0.5$",color=colors[2])
    ax.set_xlabel(r'$x/L $')
    ax.set_ylabel(r"$ U(x)/k_BT $")

    ax.plot(x_harm,(x_harm-x_min)**2/(2*a)*gamma*np.cos(x_min/a)+potential(x_min,Omega,gamma,a)-shift,linestyle=":",color='black',label=r"HA")

    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'$-0.5$', r'$0$',r'$0.5$' ,'$1.0$', r'$1.5$' ,r'$2$',  r'$2.5$', r'$3.0$']
    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"$f/u$",frameon=False)
    
        
    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,a)-shift, linestyle="-", color=colors[5])
    inset_ax.plot(x / a, potential(x, Omega, 1,a)-shift, linestyle="-", color=colors[3])
    inset_ax.plot(x / a, potential(x, Omega, gamma,a)-shift, color=colors[2])
    inset_ax.plot(x_harm, (x_harm - x_min)**2 / (2 * a) * gamma * np.cos(x_min / a) + potential(x_min, Omega, gamma,a)-shift,
                  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/potential.pdf")
    #plt.savefig('../figures/washboard_potential_inset.pdf')
    plt.show()

plot_washboard_with_HP()