# Copyright 2023 The PyGlove Authors
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Tests for pg.Ref."""

import copy
import inspect
import pickle
from typing import Any
import unittest

from pyglove.core import typing as pg_typing
from pyglove.core.symbolic import ref
from pyglove.core.symbolic.base import contains
from pyglove.core.symbolic.dict import Dict
from pyglove.core.symbolic.object import Object


class A(Object):
  x: int


class RefTest(unittest.TestCase):

  def test_basics(self):

    a = A(1)
    r = ref.Ref(a)
    self.assertIs(r.value, a)
    self.assertIs(r.infer(), a)
    self.assertEqual(ref.Ref(a), r)
    b = copy.copy(a)
    self.assertNotEqual(ref.Ref(b), r)

  def test_new(self):
    self.assertIsInstance(ref.Ref(1), int)
    self.assertIsInstance(ref.Ref(1.0), float)
    self.assertIsInstance(ref.Ref(True), bool)
    self.assertIsInstance(ref.Ref('abc'), str)

    r = ref.Ref(A(1))
    self.assertIsInstance(r, ref.Ref)
    self.assertIsInstance(r.infer(), A)
    self.assertEqual(ref.Ref(r), r)

  def test_type_check(self):

    class B:
      pass

    r = ref.Ref(A(1))
    self.assertIs(pg_typing.Object(A).apply(r), r)
    with self.assertRaisesRegex(
        TypeError, 'Expect .* but encountered .*'
    ):
      pg_typing.Object(A).apply(ref.Ref(B()))

  def test_partial_check(self):
    r = ref.Ref(A.partial())
    self.assertIs(pg_typing.Object(A).apply(r, allow_partial=True), r)
    with self.assertRaisesRegex(
        ValueError, '.* is not fully bound'
    ):
      pg_typing.Object(A).apply(r)

  def test_clone_basics(self):
    r = ref.Ref(A(1))
    self.assertIs(r.clone().infer(), r.infer())
    self.assertIs(r.clone(deep=True).infer(), r.infer())

    self.assertIs(copy.copy(r).infer(), r.infer())
    self.assertIs(copy.deepcopy(r).infer(), r.infer())

  def test_self_reference(self):

    class B(Object):
      x: Any
      y: Any

    b = B.partial(x=1)
    with self.assertRaisesRegex(
        NotImplementedError, 'Self-referential object is not supported.'
    ):
      b.rebind(y=ref.Ref(b))

  def test_repr(self):
    self.assertEqual(repr(ref.Ref(A(1))), 'Ref(A(x=1))')

  def test_str(self):
    self.assertEqual(
        str(ref.Ref(A(1))),
        inspect.cleandoc("""
        Ref(
          value = A(
            x = 1
          )
        )
        """))

  def test_to_json(self):
    with self.assertRaisesRegex(
        TypeError, '.* cannot be serialized at the moment'):
      ref.Ref(A(1)).to_json()

  def test_pickle(self):
    with self.assertRaisesRegex(
        TypeError, '.* cannot be pickled at the moment'):
      pickle.dumps(ref.Ref(A(1)))

  def test_maybe_ref(self):
    a = ref.maybe_ref(A(1))
    self.assertIsInstance(a, A)
    self.assertNotIsInstance(a, ref.Ref)

    b = Dict(x=a)
    x = ref.maybe_ref(b.x)
    self.assertIsInstance(x, ref.Ref)
    self.assertIs(x.value, a)

  def test_deref(self):
    class Foo(Object):
      x: int
      y: Any

    # Deref on non-ref value.
    a = Foo(1, 2)
    self.assertIs(ref.deref(a), a)
    self.assertEqual(ref.deref(a), Foo(1, 2))

    # Deref top-level Ref.
    a = ref.Ref(Foo(1, 2))
    self.assertIs(ref.deref(a), a.value)

    a = ref.Ref(Foo(1, ref.Ref(Foo(2, ref.Ref(Foo(3, 4))))))
    self.assertTrue(contains(a, type=ref.Ref))
    a_prime = ref.deref(a)
    self.assertIs(a_prime, a.value)
    self.assertIsInstance(a_prime.sym_getattr('y'), ref.Ref)
    self.assertIsInstance(a_prime.y.sym_getattr('y'), ref.Ref)

    # Deref entire tree.
    a = ref.Ref(Foo(1, ref.Ref(Foo(2, ref.Ref(Foo(3, 4))))))
    self.assertTrue(contains(a, type=ref.Ref))
    a_prime = ref.deref(a, recursive=True)
    self.assertIs(a_prime, a.value)
    self.assertFalse(contains(a_prime, type=ref.Ref))

  def test_to_html(self):

    def assert_style(html, expected):
      expected = inspect.cleandoc(expected).strip()
      actual = html.style_section.strip()
      if actual != expected:
        print(actual)
      self.assertEqual(actual.strip(), expected)

    def assert_content(html, expected):
      expected = inspect.cleandoc(expected).strip()
      actual = html.content.strip()
      if actual != expected:
        print(actual)
      self.assertEqual(actual.strip(), expected)

    class Foo(Object):
      x: Any

    assert_style(
        Foo(ref.Ref(Foo(1))).to_html(),
        """
        <style>
        /* Tooltip styles. */
        span.tooltip {
          visibility: hidden;
          white-space: pre-wrap;
          font-weight: normal;
          background-color: #484848;
          color: #fff;
          padding: 10px;
          border-radius: 6px;
          position: absolute;
          z-index: 1;
        }
        /* Summary styles. */
        details.pyglove summary {
          font-weight: bold;
          margin: -0.5em -0.5em 0;
          padding: 0.5em;
        }
        .summary_name {
          display: inline;
          padding: 0 5px;
        }
        .summary_title {
          display: inline;
        }
        .summary_name + div.summary_title {
          display: inline;
          color: #aaa;
        }
        .summary_title:hover + span.tooltip {
          visibility: visible;
        }
        /* Type-specific styles. */
        .pyglove.str .summary_title {
          color: darkred;
          font-style: italic;
        }
        /* Object key styles. */
        .object_key {
          margin-right: 0.25em;
        }
        .object_key:hover + .tooltip {
          visibility: visible;
          background-color: darkblue;
        }
        .object_key.str {
          color: gray;
          border: 1px solid lightgray;
          background-color: ButtonFace;
          border-radius: 0.2em;
          padding: 0.3em;
        }
        .object_key.int::before{
          content: '[';
        }
        .object_key.int::after{
          content: ']';
        }
        .object_key.int{
          border: 0;
          color: lightgray;
          background-color: transparent;
          border-radius: 0;
          padding: 0;
        }
        /* Simple value styles. */
        .simple_value {
          color: blue;
          display: inline-block;
          white-space: pre-wrap;
          padding: 0.2em;
          margin-top: 0.15em;
        }
        .simple_value.str {
          color: darkred;
          font-style: italic;
        }
        .simple_value.int, .simple_value.float {
          color: darkblue;
        }
        /* Complex value styles. */
        span.empty_container::before {
            content: '(empty)';
            font-style: italic;
            margin-left: 0.5em;
            color: #aaa;
        }
        /* Value details styles. */
        details.pyglove {
          border: 1px solid #aaa;
          border-radius: 4px;
          padding: 0.5em 0.5em 0;
          margin: 0.1em 0;
        }
        details.pyglove.special_value {
          margin-bottom: 0.75em;
        }
        details.pyglove[open] {
          padding: 0.5em 0.5em 0.5em;
        }
        .highlight {
          background-color: Mark;
        }
        .lowlight {
          opacity: 0.2;
        }
        </style>
        """
    )
    assert_content(
        Foo(ref.Ref(Foo(1))).to_html(
            use_inferred=False,
            enable_summary_tooltip=False,
            enable_key_tooltip=False,
        ),
        """
        <details open class="pyglove foo"><summary><div class="summary_title">Foo(...)</div></summary><div class="complex_value foo"><table><tr><td><span class="object_key str">x</span></td><td><div><details class="pyglove ref"><summary><div class="summary_title">Foo(...)</div></summary><div class="complex_value foo"><table><tr><td><span class="object_key str">x</span></td><td><div><span class="simple_value int">1</span></div></td></tr></table></div></details></div></td></tr></table></div></details>
        """
    )


if __name__ == '__main__':
  unittest.main()
