import openturns as ot
import openturns.experimental as otexp
from matplotlib import pyplot as plt
import openturns.viewer as otv

ot.RandomGenerator.SetSeed(0)
if hasattr(ot, "FiniteDiscreteDistributionFactory"):
    factory = ot.FiniteDiscreteDistributionFactory()
else:
    factory = otexp.FiniteDiscreteDistributionFactory()
ref = factory.build()
dimension = ref.getDimension()
if dimension <= 2:
    if dimension == 1:
        sample = ref.getSample(50)
        distribution = factory.build(sample)
        distribution.setDescription(['$t$'])
        pdf_graph = distribution.drawPDF(256)
        cloud = ot.Cloud(sample, ot.Sample(sample.getSize(), 1))
        cloud.setColor('blue')
        cloud.setPointStyle('fcircle')
        pdf_graph.add(cloud)
        pdf_graph.setTitle(str(distribution))
        fig = plt.figure(figsize=(10, 4))
        pdf_axis = fig.add_subplot(111)
        otv.View(pdf_graph, figure=fig, axes=[pdf_axis], add_legend=False)
    else:
        sample = ref.getSample(500)
        distribution = factory.build(sample)
        distribution.setDescription(['$t_0$', '$t_1$'])
        pdf_graph = distribution.drawPDF([256]*2)
        cloud = ot.Cloud(sample)
        cloud.setColor('red')
        cloud.setPointStyle('fcircle')
        pdf_graph.add(cloud)
        pdf_graph.setTitle(str(distribution))
        fig = plt.figure(figsize=(10, 4))
        pdf_axis = fig.add_subplot(111)
        otv.View(pdf_graph, figure=fig, axes=[pdf_axis], add_legend=False, square_axes=True)