from math import sqrt

from bokeh.models import Arrow, Label
from bokeh.plotting import figure, output_file, show
from bokeh.util.compiler import TypeScript

TS_CODE = """
import * as p from "core/properties"
import {Label, LabelView} from "models/annotations/label"
declare const katex: any

export class LatexLabelView extends LabelView {
  model: LatexLabel

  render(): void {
    //--- Start of copied section from ``Label.render`` implementation

    // Here because AngleSpec does units tranform and label doesn't support specs
    let angle: number
    switch (this.model.angle_units) {
      case "rad": {
        angle = -this.model.angle
        break
      }
      case "deg": {
        angle = (-this.model.angle * Math.PI) / 180.0
        break
      }
      default:
        throw new Error("unreachable code")
    }

    const panel = this.panel != null ? this.panel : this.plot_view.frame

    const xscale = this.plot_view.frame.xscales[this.model.x_range_name]
    const yscale = this.plot_view.frame.yscales[this.model.y_range_name]

    let sx = this.model.x_units == "data" ? xscale.compute(this.model.x) : panel.xview.compute(this.model.x)
    let sy = this.model.y_units == "data" ? yscale.compute(this.model.y) : panel.yview.compute(this.model.y)

    sx += this.model.x_offset
    sy -= this.model.y_offset

    //--- End of copied section from ``Label.render`` implementation
    // Must render as superpositioned div (not on canvas) so that KaTex
    // css can properly style the text
    this._css_text(this.plot_view.canvas_view.ctx, "", sx, sy, angle)

    // ``katex`` is loaded into the global window at runtime
    // katex.renderToString returns a html ``span`` element
    katex.render(this.model.text, this.el, {displayMode: true})
  }
}

export namespace LatexLabel {
  export type Attrs = p.AttrsOf<Props>

  export type Props = Label.Props
}

export interface LatexLabel extends LatexLabel.Attrs {}

export class LatexLabel extends Label {
  properties: LatexLabel.Props

  constructor(attrs?: Partial<LatexLabel.Attrs>) {
    super(attrs)
  }

  static init_LatexLabel() {
    this.prototype.default_view = LatexLabelView
  }
}
"""


class LatexLabel(Label):
    """A subclass of the Bokeh built-in `Label` that supports rendering
    LaTex using the KaTex typesetting library.

    Only the render method of LabelView is overloaded to perform the
    text -> latex (via katex) conversion. Note: ``render_mode="canvas``
    isn't supported and certain DOM manipulation happens in the Label
    superclass implementation that requires explicitly setting
    `render_mode='css'`).
    """

    __javascript__ = ["https://cdnjs.cloudflare.com/ajax/libs/KaTeX/0.6.0/katex.min.js"]
    __css__ = ["https://cdnjs.cloudflare.com/ajax/libs/KaTeX/0.6.0/katex.min.css"]
    __implementation__ = TypeScript(TS_CODE)


class PlaneStress:
    def __init__(self, sigma_x, sigma_y, tau_xy):
        self.sigma_x = sigma_x
        self.sigma_y = sigma_y
        self.tau_xy = tau_xy

        self.tau_max = sqrt(((self.sigma_x - self.sigma_y) / 2) ** 2 + self.tau_xy ** 2)
        self.sigma_1 = ((self.sigma_x + self.sigma_y) / 2) + self.tau_max
        self.sigma_2 = ((self.sigma_x + self.sigma_y) / 2) - self.tau_max

    def plane(self):
        plot = figure(x_range=(-5, 5), y_range=(-5, 5))

        plot.axis.major_label_text_font_size = "0pt"
        plot.axis.major_tick_line_color = None
        plot.axis[0].ticker.num_minor_ticks = 0
        plot.axis[1].ticker.num_minor_ticks = 0
        plot.grid.visible = False
        plot.axis.visible = False

        plot.rect(0, 0, 4, 4, fill_alpha=0, line_color="black", line_width=3)

        # sigma_x
        plot.add_layout(Arrow(x_start=2, y_start=0, x_end=4.5, y_end=0))
        plot.add_layout(Arrow(x_start=-2, y_start=0, x_end=-4.5, y_end=0))
        plot.add_layout(
            LatexLabel(
                x=3.5, y=0.75, text=f"\\sigma_{{x}} = {self.sigma_x}", render_mode="css"
            )
        )

        # sigma_y
        plot.add_layout(Arrow(x_start=0, y_start=2, x_end=0, y_end=4.5))
        plot.add_layout(Arrow(x_start=0, y_start=-2, x_end=0, y_end=-4.5))
        plot.add_layout(
            LatexLabel(
                x=0.25, y=4, text=f"\\sigma_{{y}} = {self.sigma_y}", render_mode="css"
            )
        )

        # tau_xy
        plot.add_layout(Arrow(x_start=2.5, y_start=-2, x_end=2.5, y_end=2.2))
        plot.add_layout(Arrow(y_start=2.5, x_start=-2, y_end=2.5, x_end=2.2))
        plot.add_layout(
            LatexLabel(
                x=2.5, y=2.6, text=f"\\tau_{{xy}} = {self.tau_xy}", render_mode="css"
            )
        )
        plot.add_layout(Arrow(x_start=-2.5, y_start=2, x_end=-2.5, y_end=-2.2))
        plot.add_layout(Arrow(y_start=-2.5, x_start=2, y_end=-2.5, x_end=-2.2))

        show(plot)


x = PlaneStress(80, -40, 25)

x.plane()