{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "using Plots, LinearAlgebra\n",
    "include(\"p_space.jl\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "vertices = (0, 0.4 - 0.2im, 0.8, 6)\n",
    "subdivisions = 128\n",
    "p, w = get_mesh(vertices, subdivisions)\n",
    "mesh_E = p.*p ./ (2*0.5)\n",
    "\n",
    "# ResonanceEC: Eq. (20)\n",
    "V_system(c) = (p, q) -> c*(-5*g0(sqrt(3), p, q) + 2*g0(sqrt(10), p, q))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "training_points = range(0.75, 0.45, 5)\n",
    "training_E = Vector{ComplexF64}(undef, length(training_points))\n",
    "EC_basis = Matrix{ComplexF64}(undef, length(p), length(training_points))\n",
    "\n",
    "for (j, c) in enumerate(training_points)\n",
    "    evals, evecs = eigen(get_H_matrix(V_system(c), p, w))\n",
    "    i = identify_pole_i(p, evals)\n",
    "    training_E[j] = evals[i]\n",
    "    EC_basis[:, j] = evecs[:, i]\n",
    "end\n",
    "\n",
    "scatter(real.(training_E), imag.(training_E), label=\"training\")\n",
    "plot!(real.(mesh_E), imag.(mesh_E), label=\"contour\")\n",
    "xlims!(0,1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "extrapolate_points = range(0.40, 0.20, 5)\n",
    "ref_E = 0.2 - 0.1im\n",
    "\n",
    "exact_E = Vector{ComplexF64}(undef, length(extrapolate_points))\n",
    "extrapolate_E = Vector{ComplexF64}(undef, length(extrapolate_points))\n",
    "\n",
    "for (j, c) in enumerate(extrapolate_points)\n",
    "    exact_E[j] = quick_pole_E(V_system(c))\n",
    "\n",
    "    EC_basis_w = EC_basis .* w\n",
    "    H = get_H_matrix(V_system(c), p, w)\n",
    "    H_EC = transpose(EC_basis) * H * EC_basis\n",
    "    N_EC = transpose(EC_basis) * EC_basis\n",
    "    evals = eigvals(H_EC, N_EC)\n",
    "    i = argmin(abs.(evals .- ref_E))\n",
    "    ref_E = evals[i]\n",
    "    extrapolate_E[j] = evals[i]\n",
    "end\n",
    "\n",
    "scatter(real.(training_E), imag.(training_E), label=\"training\")\n",
    "scatter!(real.(exact_E), imag.(exact_E), label=\"exact\")\n",
    "scatter!(real.(extrapolate_E), imag.(extrapolate_E), label=\"extrapolated\")\n",
    "plot!(real.(mesh_E), imag.(mesh_E), label=\"contour\")\n",
    "xlims!(0,1)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Julia 1.9.0",
   "language": "julia",
   "name": "julia-1.9"
  },
  "language_info": {
   "file_extension": ".jl",
   "mimetype": "application/julia",
   "name": "julia",
   "version": "1.9.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}