# -*- coding: utf-8 -*-
# ____________________________________________________________________________________
#
# Pyomo: Python Optimization Modeling Objects
# Copyright (c) 2008-2026 National Technology and Engineering Solutions of Sandia, LLC
# Under the terms of Contract DE-NA0003525 with National Technology and Engineering
# Solutions of Sandia, LLC, the U.S. Government retains certain rights in this
# software.  This software is distributed under the 3-clause BSD License.
# ____________________________________________________________________________________
from collections import namedtuple
from io import StringIO

import pyomo.common.unittest as unittest

from pyomo.common.formatting import tostr, tabular_writer, StreamIndenter


class DerivedList(list):
    pass


class DerivedTuple(tuple):
    pass


class DerivedDict(dict):
    pass


class DerivedStr(str):
    pass


NamedTuple = namedtuple('NamedTuple', ['x', 'y'])


class TestToStr(unittest.TestCase):
    def test_new_type_float(self):
        self.assertEqual(tostr(0.5), '0.5')
        self.assertIs(tostr.handlers[float], tostr.handlers[None])

    def test_new_type_int(self):
        self.assertEqual(tostr(0), '0')
        self.assertIs(tostr.handlers[int], tostr.handlers[None])

    def test_new_type_str(self):
        self.assertEqual(tostr(DerivedStr(1)), '1')
        self.assertIs(tostr.handlers[DerivedStr], tostr.handlers[str])

    def test_new_type_list(self):
        self.assertEqual(tostr(DerivedList([1, 2])), '[1, 2]')
        self.assertIs(tostr.handlers[DerivedList], tostr.handlers[list])

    def test_new_type_dict(self):
        self.assertEqual(tostr(DerivedDict({1: 2})), '{1: 2}')
        self.assertIs(tostr.handlers[DerivedDict], tostr.handlers[dict])

    def test_new_type_tuple(self):
        self.assertEqual(tostr(DerivedTuple([1, 2])), '(1, 2)')
        self.assertIs(tostr.handlers[DerivedTuple], tostr.handlers[tuple])

    def test_new_type_namedtuple(self):
        self.assertEqual(tostr(NamedTuple(1, 2)), 'NamedTuple(x=1, y=2)')
        self.assertIs(tostr.handlers[NamedTuple], tostr.handlers[None])


class TestTabularWriter(unittest.TestCase):
    def test_unicode_table(self):
        # Test that an embedded unicode character does not foul up the
        # table alignment
        os = StringIO()
        data = {1: ("a", 1), (2, 3): ("∧", 2)}
        tabular_writer(os, "", data.items(), ["s", "val"], lambda k, v: v)
        ref = u"""
Key    : s : val
     1 : a :   1
(2, 3) : ∧ :   2
"""
        self.assertEqual(ref.strip(), os.getvalue().strip())

    def test_tuple_list_dict(self):
        os = StringIO()
        data = {(1,): (["a", 1], 1), ('2', 3): ({1: 'a', 2: '2'}, '2')}
        tabular_writer(os, "", data.items(), ["s", "val"], lambda k, v: v)
        ref = u"""
Key      : s                : val
    (1,) :         ['a', 1] :   1
('2', 3) : {1: 'a', 2: '2'} :   2
"""
        self.assertEqual(ref.strip(), os.getvalue().strip())

    def test_no_header(self):
        os = StringIO()
        data = {(2,): (["a", 1], 1), (1, 3): ({1: 'a', 2: '2'}, '2')}
        tabular_writer(os, "", data.items(), [], lambda k, v: v)
        ref = u"""
        ['a', 1] : 1
{1: 'a', 2: '2'} : 2
"""
        self.assertEqual(ref.strip(), os.getvalue().strip())

    def test_no_data(self):
        os = StringIO()
        data = {}
        tabular_writer(os, "", data.items(), ['s', 'val'], lambda k, v: v)
        ref = u"""
Key : s : val
"""
        self.assertEqual(ref.strip(), os.getvalue().strip())

    def test_multiline_generator(self):
        os = StringIO()
        data = {'a': 0, 'b': 1, 'c': 3}

        def _data_gen(i, j):
            for n in range(j):
                yield (n, chr(ord('a') + n) * j)

        tabular_writer(os, "", data.items(), ['i', 'j'], _data_gen)
        ref = u"""
Key : i    : j
  a : None : None
  b :    0 :    a
  c :    0 :  aaa
    :    1 :  bbb
    :    2 :  ccc
"""
        self.assertEqual(ref.strip(), os.getvalue().strip())

    def test_multiline_generator_exception(self):
        os = StringIO()
        data = {'a': 0, 'b': 1, 'c': 3}

        def _data_gen(i, j):
            if i == 'b':
                raise ValueError("invalid")
            for n in range(j):
                yield (n, chr(ord('a') + n) * j)

        tabular_writer(os, "", data.items(), ['i', 'j'], _data_gen)
        ref = u"""
Key : i    : j
  a : None : None
  b : None : None
  c :    0 :  aaa
    :    1 :  bbb
    :    2 :  ccc
"""
        self.assertEqual(ref.strip(), os.getvalue().strip())

    def test_data_exception(self):
        os = StringIO()
        data = {'a': 0, 'b': 1, 'c': 3}

        def _data_gen(i, j):
            if i == 'b':
                raise ValueError("invalid")
            return (j, i * (j + 1))

        tabular_writer(os, "", data.items(), ['i', 'j'], _data_gen)
        ref = u"""
Key : i    : j
  a :    0 :    a
  b : None : None
  c :    3 : cccc
"""
        self.assertEqual(ref.strip(), os.getvalue().strip())

    def test_multiline_alignment(self):
        os = StringIO()
        data = {'a': 1, 'b': 2, 'c': 3}

        def _data_gen(i, j):
            for n in range(j):
                _str = chr(ord('a') + n) * (j + 1)
                if n % 2:
                    _str = list(_str)
                    _str[1] = ' '
                    _str = ''.join(_str)
                yield (n, _str)

        tabular_writer(os, "", data.items(), ['i', 'j'], _data_gen)
        ref = u"""
Key : i : j
  a : 0 : aa
  b : 0 : aaa
    : 1 : b b
  c : 0 : aaaa
    : 1 : b bb
    : 2 : cccc
"""
        self.assertEqual(ref.strip(), os.getvalue().strip())


class TestStreamIndenter(unittest.TestCase):
    def test_empty(self):
        OUT1 = StringIO()
        OUT2 = StreamIndenter(OUT1)
        self.assertEqual(0, OUT2.write(''))
        self.assertEqual('', OUT2.getvalue())

    def test_noprefix(self):
        OUT1 = StringIO()
        OUT2 = StreamIndenter(OUT1)
        self.assertEqual(28, OUT2.write('Hello?\nHello, world!'))
        self.assertEqual('    Hello?\n    Hello, world!', OUT2.getvalue())

    def test_prefix(self):
        prefix = 'foo: '
        OUT1 = StringIO()
        OUT2 = StreamIndenter(OUT1, prefix)
        OUT2.write('Hello?\nText\n\nHello, world!')
        self.assertEqual(
            'foo: Hello?\nfoo: Text\nfoo:\nfoo: Hello, world!', OUT2.getvalue()
        )

    def test_blank_lines(self):
        OUT1 = StringIO()
        OUT2 = StreamIndenter(OUT1)
        OUT2.write('Hello?\n\nText\n\nHello, world!')
        self.assertEqual('    Hello?\n\n    Text\n\n    Hello, world!', OUT2.getvalue())

    def test_blank_lines_nonwhitespace_indent(self):
        OUT1 = StringIO()
        OUT2 = StreamIndenter(OUT1, " | ")
        OUT2.write('Hello?\n\nText\n')
        OUT2.write('\n')
        OUT2.write('Hello, world!')
        self.assertEqual(
            ' | Hello?\n |\n | Text\n |\n | Hello, world!', OUT2.getvalue()
        )

    def test_writelines(self):
        OUT1 = StringIO()
        OUT2 = StreamIndenter(OUT1)
        OUT2.writelines(['Hello?\n', '\n', 'Text\n', '\n', 'Hello, world!'])
        self.assertEqual('    Hello?\n\n    Text\n\n    Hello, world!', OUT2.getvalue())

    def test_nested(self):
        OUT1 = StringIO()
        OUT2 = StreamIndenter(OUT1)
        OUT3 = StreamIndenter(OUT2)
        self.assertIs(OUT3.target_os, OUT2.target_os)
        self.assertIs(OUT3.target_os, OUT1)
        OUT3.write('Hello?\n\nText\n\nHello, world!')
        self.assertEqual(
            '        Hello?\n\n        Text\n\n        Hello, world!', OUT1.getvalue()
        )

    def test_nested_interleave(self):
        OUT1 = StringIO()
        OUT2 = StreamIndenter(OUT1)
        OUT3 = StreamIndenter(OUT2)
        self.assertIs(OUT3.target_os, OUT2.target_os)
        self.assertIs(OUT3.target_os, OUT1)
        OUT3.write('Hello?')
        OUT2.write('\n\n')
        OUT3.write('Text\n')
        OUT2.write('Hi\n')
        OUT3.write('Hello, world!')
        self.assertEqual(
            '        Hello?\n\n        Text\n    Hi\n        Hello, world!',
            OUT1.getvalue(),
        )
