import os
import stat
import string
from absl.testing import absltest
from src.test.py.bazel import test_base

# pylint: disable=g-import-not-at-top
if os.name == 'nt':
  import win32api

class LauncherTest(test_base.TestBase):

  def _buildJavaTargets(self, bazel_bin, binary_suffix):
    self.RunBazel(['build', '//foo'])
    main_binary = os.path.join(bazel_bin, 'foo/foo%s' % binary_suffix)
    self.assertTrue(os.path.isfile(main_binary))
    self.assertTrue(
        os.path.isdir(
            os.path.join(bazel_bin, 'foo/foo%s.runfiles' % binary_suffix)))

    if self.IsWindows():
      self.assertTrue(os.path.isfile(main_binary))
      self.AssertRunfilesManifestContains(
          os.path.join(
              bazel_bin, 'foo/foo%s.runfiles/MANIFEST' % binary_suffix
          ),
          '_main/bar/bar.txt',
      )
    else:
      self.assertTrue(
          os.path.islink(
              os.path.join(bazel_bin, 'foo/foo.runfiles/_main/bar/bar.txt')
          )
      )

    _, stdout, _ = self.RunProgram([main_binary])
    self.assertEqual(len(stdout), 4)
    self.assertEqual(stdout[0], 'hello java')
    if self.IsWindows():
      self.assertRegexpMatches(
          stdout[1], r'java_runfiles=.*foo\\foo%s.runfiles' % binary_suffix)
      self.assertEqual(stdout[2], 'runfiles_manifest_only=1')
      self.assertRegexpMatches(
          stdout[3], r'^runfiles_manifest_file=[a-zA-Z]:[/\\].*MANIFEST$')
    else:
      self.assertRegexpMatches(stdout[1], r'java_runfiles=.*/foo/foo.runfiles')
      self.assertEqual(stdout[2], 'runfiles_manifest_only=')
      self.assertRegexpMatches(stdout[3], r'^runfiles_manifest_file.*MANIFEST$')

  def _buildShBinaryTargets(self, bazel_bin, bin1_suffix):
    self.RunBazel(['build', '//foo:bin1.sh'])

    bin1 = os.path.join(bazel_bin, 'foo', 'bin1.sh%s' % bin1_suffix)

    self.assertTrue(os.path.exists(bin1))
    self.assertTrue(
        os.path.isdir(
            os.path.join(bazel_bin, 'foo/bin1.sh%s.runfiles' % bin1_suffix)))

    self.RunBazel(['build', '//foo:bin2.cmd'])

    bin2 = os.path.join(bazel_bin, 'foo/bin2.cmd')
    self.assertTrue(os.path.exists(bin2))
    self.assertTrue(
        os.path.isdir(os.path.join(bazel_bin, 'foo/bin2.cmd.runfiles')))

    exit_code, _, stderr = self.RunBazel(
        ['build', '//foo:bin3.bat'], allow_failure=True
    )
    if self.IsWindows():
      self.AssertExitCode(exit_code, 1, stderr)
      self.assertIn('target name extension should match source file extension',
                    os.linesep.join(stderr))
    else:
      bin3 = os.path.join(bazel_bin, 'foo', 'bin3.bat')
      self.assertTrue(os.path.exists(bin3))
      self.assertTrue(
          os.path.isdir(os.path.join(bazel_bin, 'foo/bin3.bat.runfiles')))

    if self.IsWindows():
      self.assertTrue(os.path.isfile(bin1))
      self.assertTrue(os.path.isfile(bin2))
    else:
      self.assertTrue(os.path.islink(bin1))
      self.assertTrue(os.path.islink(bin2))
      self.assertTrue(os.path.islink(bin3))

    if self.IsWindows():
      self.AssertRunfilesManifestContains(
          os.path.join(
              bazel_bin, 'foo/bin1.sh%s.runfiles/MANIFEST' % bin1_suffix
          ),
          '_main/bar/bar.txt',
      )
      self.AssertRunfilesManifestContains(
          os.path.join(bazel_bin, 'foo/bin2.cmd.runfiles/MANIFEST'),
          '_main/bar/bar.txt',
      )
    else:
      self.assertTrue(
          os.path.islink(
              os.path.join(bazel_bin, 'foo/bin1.sh.runfiles/_main/bar/bar.txt')
          )
      )
      self.assertTrue(
          os.path.islink(
              os.path.join(bazel_bin, 'foo/bin2.cmd.runfiles/_main/bar/bar.txt')
          )
      )
      self.assertTrue(
          os.path.islink(
              os.path.join(bazel_bin, 'foo/bin3.bat.runfiles/_main/bar/bar.txt')
          )
      )

    _, stdout, _ = self.RunProgram([bin1])
    self.assertEqual(len(stdout), 3)
    self.assertEqual(stdout[0], 'hello shell')
    if self.IsWindows():
      self.assertEqual(stdout[1], 'runfiles_manifest_only=1')
      self.assertRegexpMatches(
          stdout[2],
          (r'^runfiles_manifest_file='
           r'[a-zA-Z]:/.*/foo/bin1.sh%s.runfiles/MANIFEST$' % bin1_suffix))
    else:
      # TODO(laszlocsomor): Find out whether the runfiles-related envvars should
      # be set on Linux (e.g. $RUNFILES, $RUNFILES_MANIFEST_FILE). Currently
      # they aren't, and that may be a bug. If it's indeed a bug, fix that bug
      # and update this test.
      self.assertEqual(stdout[1], 'runfiles_manifest_only=')
      self.assertEqual(stdout[2], 'runfiles_manifest_file=')

    if self.IsWindows():
      exit_code, stdout, stderr = self.RunProgram([bin2])
      self.AssertExitCode(exit_code, 0, stderr)
      self.assertEqual(stdout[0], 'hello batch')

  def _buildPyTargets(self, bazel_bin, binary_suffix):
    # Verify that the build of our py_binary succeeds.
    self.RunBazel(['build', '//foo:foo'])

    # Verify that generated files exist.
    foo_bin = os.path.join(bazel_bin, 'foo', 'foo%s' % binary_suffix)
    self.assertTrue(os.path.isfile(foo_bin))
    self.assertTrue(
        os.path.isdir(
            os.path.join(bazel_bin, 'foo/foo%s.runfiles' % binary_suffix)))

    # Verify contents of runfiles (manifest).
    if self.IsWindows():
      self.AssertRunfilesManifestContains(
          os.path.join(
              bazel_bin, 'foo/foo%s.runfiles/MANIFEST' % binary_suffix
          ),
          '_main/bar/bar.txt',
      )
    else:
      self.assertTrue(
          os.path.islink(
              os.path.join(bazel_bin, 'foo/foo.runfiles/_main/bar/bar.txt')
          )
      )

    # Try to run the built py_binary.
    _, stdout, _ = self.RunProgram([foo_bin])
    self.assertEqual(stdout[0], 'Hello World!')

    # Try to use the py_binary as an executable in a Starlark rule.
    self.RunBazel(['build', '//foo:hello'])

    # Verify that the Starlark action generated the right output.
    hello_path = os.path.join(bazel_bin, 'foo', 'hello.txt')
    self.assertTrue(os.path.isfile(hello_path))
    with open(hello_path, 'r') as f:
      self.assertEqual(f.read(), 'Hello World!')

    # Verify that running py_test succeeds.
    self.RunBazel(['test', '//foo:test'])

  def _buildAndCheckArgumentPassing(self, package, target_name):
    _, stdout, _ = self.RunBazel(['info', 'bazel-bin'])
    bazel_bin = stdout[0]

    self.RunBazel(['build', '//%s:%s' % (package, target_name)])

    bin_suffix = '.exe' if self.IsWindows() else ''
    bin1 = os.path.join(bazel_bin, package, '%s%s' % (target_name, bin_suffix))
    self.assertTrue(os.path.exists(bin1))
    self.assertTrue(
        os.path.isdir(
            os.path.join(bazel_bin, '%s/%s%s.runfiles' % (package, target_name,
                                                          bin_suffix))))

    arguments = ['a', 'a b', '"b"', 'C:\\a\\b\\', '"C:\\a b\\c\\"']
    _, stdout, _ = self.RunProgram([bin1] + arguments)
    self.assertEqual(stdout, arguments)

  def testJavaBinaryLauncher(self):
    self.CreateWorkspaceWithDefaultRepos('WORKSPACE')
    self.ScratchFile('foo/BUILD', [
        'java_binary(',
        '  name = "foo",',
        '  srcs = ["Main.java"],',
        '  main_class = "Main",',
        '  data = ["//bar:bar.txt"],',
        ')',
    ])
    self.ScratchFile('foo/Main.java', [
        'public class Main {',
        '  public static void main(String[] args) {'
        '    System.out.println("hello java");',
        '    System.out.println("java_runfiles=" + ',
        '        System.getenv("JAVA_RUNFILES"));',
        '    System.out.println("runfiles_manifest_only=" + ',
        '        System.getenv("RUNFILES_MANIFEST_ONLY"));',
        '    System.out.println("runfiles_manifest_file=" + ',
        '        System.getenv("RUNFILES_MANIFEST_FILE"));',
        '  }',
        '}',
    ])
    self.ScratchFile('bar/BUILD', ['exports_files(["bar.txt"])'])
    self.ScratchFile('bar/bar.txt', ['hello'])

    _, stdout, _ = self.RunBazel(['info', 'bazel-bin'])
    bazel_bin = stdout[0]
    self._buildJavaTargets(bazel_bin, '.exe' if self.IsWindows() else '')

  def testJavaBinaryArgumentPassing(self):
    self.CreateWorkspaceWithDefaultRepos('WORKSPACE')
    self.ScratchFile('foo/BUILD', [
        'java_binary(',
        '  name = "bin",',
        '  srcs = ["Main.java"],',
        '  main_class = "Main",',
        ')',
    ])
    self.ScratchFile('foo/Main.java', [
        'public class Main {',
        '  public static void main(String[] args) {'
        '    for (String arg : args) {',
        '      System.out.println(arg);',
        '    }'
        '  }',
        '}',
    ])

    self._buildAndCheckArgumentPassing('foo', 'bin')

  def testShBinaryLauncher(self):
    self.CreateWorkspaceWithDefaultRepos('WORKSPACE')
    self.ScratchFile(
        'foo/BUILD',
        [
            # On Linux/MacOS, all sh_binary rules generate an output file with
            # the same name as the rule, and this is a symlink to the file in
            # `srcs`. (Bazel allows only one file in `sh_binary.srcs`.)
            # On Windows, if the srcs's extension is one of ".exe", ".cmd", or
            # ".bat", then Bazel requires the rule's name has the same
            # extension, and the output file will be a copy of the source file.
            'sh_binary(',
            '  name = "bin1.sh",',
            '  srcs = ["foo.sh"],',
            '  data = ["//bar:bar.txt"],',
            ')',
            'sh_binary(',
            '  name = "bin2.cmd",',  # name's extension matches that of srcs[0]
            '  srcs = ["foo.cmd"],',
            '  data = ["//bar:bar.txt"],',
            ')',
            'sh_binary(',
            '  name = "bin3.bat",',  # name's extension doesn't match srcs[0]'s
            '  srcs = ["foo.cmd"],',
            '  data = ["//bar:bar.txt"],',
            ')',
        ])
    foo_sh = self.ScratchFile('foo/foo.sh', [
        '#!/bin/bash',
        'echo hello shell',
        'echo runfiles_manifest_only=${RUNFILES_MANIFEST_ONLY:-}',
        'echo runfiles_manifest_file=${RUNFILES_MANIFEST_FILE:-}',
    ])
    foo_cmd = self.ScratchFile('foo/foo.cmd', ['@echo hello batch'])
    self.ScratchFile('bar/BUILD', ['exports_files(["bar.txt"])'])
    self.ScratchFile('bar/bar.txt', ['hello'])
    os.chmod(foo_sh, stat.S_IRUSR | stat.S_IWUSR | stat.S_IXUSR)
    os.chmod(foo_cmd, stat.S_IRUSR | stat.S_IWUSR | stat.S_IXUSR)

    _, stdout, _ = self.RunBazel(['info', 'bazel-bin'])
    bazel_bin = stdout[0]
    self._buildShBinaryTargets(bazel_bin, '.exe' if self.IsWindows() else '')

  def testShBinaryArgumentPassing(self):
    self.CreateWorkspaceWithDefaultRepos('WORKSPACE')
    self.ScratchFile('foo/BUILD', [
        'sh_binary(',
        '  name = "bin",',
        '  srcs = ["bin.sh"],',
        ')',
    ])
    foo_sh = self.ScratchFile('foo/bin.sh', [
        '#!/bin/bash',
        '# Store arguments in a array',
        'args=("$@")',
        '# Get the number of arguments',
        'N=${#args[@]}',
        '# Echo each argument',
        'for (( i=0;i<$N;i++)); do',
        ' echo ${args[${i}]}',
        'done',
    ])
    os.chmod(foo_sh, stat.S_IRUSR | stat.S_IWUSR | stat.S_IXUSR)

    self._buildAndCheckArgumentPassing('foo', 'bin')

  def testPyBinaryLauncher(self):
    self.CreateWorkspaceWithDefaultRepos('WORKSPACE')
    self.ScratchFile(
        'foo/foo.bzl',
        [
            'def _impl(ctx):',
            '  ctx.actions.run(',
            '      arguments=[ctx.outputs.out.path],',
            '      outputs=[ctx.outputs.out],',
            '      executable=ctx.executable._hello_world,',
            '      use_default_shell_env=True)',
            '',
            'helloworld = rule(',
            '  implementation=_impl,',
            '  attrs={',
            '      "srcs": attr.label_list(allow_files=True),',
            '      "out": attr.output(mandatory=True),',
            '      "_hello_world": attr.label(executable=True, cfg="exec",',
            '                                 allow_files=True,',
            '                                 default=Label("//foo:foo"))',
            '  }',
            ')',
        ],
    )
    self.ScratchFile('foo/BUILD', [
        'load(":foo.bzl", "helloworld")', '', 'py_binary(', '  name = "foo",',
        '  srcs = ["foo.py"],', '  data = ["//bar:bar.txt"],', ')', '',
        'py_test(', '  name = "test",', '  srcs = ["test.py"],', ')', '',
        'helloworld(', '  name = "hello",', '  out = "hello.txt",', ')'
    ])
    foo_py = self.ScratchFile('foo/foo.py', [
        '#!/usr/bin/env python3',
        'import sys',
        'if len(sys.argv) == 2:',
        '  with open(sys.argv[1], "w") as f:',
        '    f.write("Hello World!")',
        'else:',
        '  print("Hello World!")',
    ])
    test_py = self.ScratchFile('foo/test.py', [
        '#!/usr/bin/env python3',
        'import unittest',
        'class MyTest(unittest.TestCase):',
        '  def test_dummy(self):',
        '      pass',
        'if __name__ == \'__main__\':',
        '  unittest.main()',
    ])
    self.ScratchFile('bar/BUILD', ['exports_files(["bar.txt"])'])
    self.ScratchFile('bar/bar.txt', ['hello'])
    os.chmod(foo_py, stat.S_IRUSR | stat.S_IWUSR | stat.S_IXUSR)
    os.chmod(test_py, stat.S_IRUSR | stat.S_IWUSR | stat.S_IXUSR)

    _, stdout, _ = self.RunBazel(['info', 'bazel-bin'])
    bazel_bin = stdout[0]
    self._buildPyTargets(bazel_bin, '.exe' if self.IsWindows() else '')

  def testPyBinaryArgumentPassing(self):
    self.CreateWorkspaceWithDefaultRepos('WORKSPACE')
    self.ScratchFile('foo/BUILD', [
        'py_binary(',
        '  name = "bin",',
        '  srcs = ["bin.py"],',
        ')',
    ])
    self.ScratchFile('foo/bin.py', [
        'import sys',
        'for arg in sys.argv[1:]:',
        '  print(arg)',
    ])

    self._buildAndCheckArgumentPassing('foo', 'bin')

  def testPyBinaryLauncherWithDifferentArgv0(self):
    """Test for https://github.com/bazelbuild/bazel/issues/14343."""
    self.CreateWorkspaceWithDefaultRepos('WORKSPACE')
    self.ScratchFile('foo/BUILD', [
        'py_binary(',
        '  name = "bin",',
        '  srcs = ["bin.py"],',
        ')',
    ])
    self.ScratchFile('foo/bin.py', ['print("Hello world")'])

    _, stdout, _ = self.RunBazel(['info', 'bazel-bin'])
    bazel_bin = stdout[0]

    # Verify that the build of our py_binary succeeds.
    self.RunBazel(['build', '//foo:bin'])

    # Try to run the built py_binary.
    binary_suffix = '.exe' if self.IsWindows() else ''
    foo_bin = os.path.join(bazel_bin, 'foo', 'bin%s' % binary_suffix)
    args = [r'C:\Invalid.exe' if self.IsWindows() else '/invalid']
    _, stdout, _ = self.RunProgram(args, executable=foo_bin)
    self.assertEqual(stdout[0], 'Hello world')

  def testWindowsJavaExeLauncher(self):
    # Skip this test on non-Windows platforms
    if not self.IsWindows():
      return
    self.CreateWorkspaceWithDefaultRepos('WORKSPACE')
    self.ScratchFile('foo/BUILD', [
        'java_binary(',
        '  name = "foo",',
        '  srcs = ["Main.java"],',
        '  main_class = "Main",',
        '  jvm_flags = ["--flag1", "--flag2"],',
        '  data = ["advice-1.jar", "advice-2.jar"],',
        ')',
    ])
    self.ScratchFile('foo/advice-1.jar')
    self.ScratchFile('foo/advice-2.jar')
    self.ScratchFile('foo/Main.java', [
        'public class Main {',
        '  public static void main(String[] args) {',
        '    System.out.println("helloworld");',
        '  }',
        '}',
    ])

    _, stdout, _ = self.RunBazel(['info', 'bazel-bin'])
    bazel_bin = stdout[0]

    self.RunBazel(['build', '//foo:foo'])

    binary = os.path.join(bazel_bin, 'foo', 'foo.exe')
    self.assertTrue(os.path.exists(binary))

    # Add this flag to make launcher print the command it generated instead of
    # launching the real program.
    print_cmd = '--print_launcher_command'

    _, stdout, _ = self.RunProgram([binary, '--debug', print_cmd])
    self.assertIn(
        '-agentlib:jdwp=transport=dt_socket,server=y,suspend=y,address=5005',
        stdout)

    _, stdout, _ = self.RunProgram(
        [binary, '--debug', print_cmd],
        env_add={'DEFAULT_JVM_DEBUG_PORT': '12345'},
    )
    self.assertIn(
        '-agentlib:jdwp=transport=dt_socket,server=y,suspend=y,address=12345',
        stdout)

    _, stdout, _ = self.RunProgram(
        [binary, '--debug=12345', print_cmd],
        env_add={
            'DEFAULT_JVM_DEBUG_SUSPEND': 'n',
            'PERSISTENT_TEST_RUNNER': 'true',
        },
    )
    self.assertIn(
        '-agentlib:jdwp=transport=dt_socket,server=y,suspend=n,address=12345'
        ',quiet=y', stdout)

    _, stdout, _ = self.RunProgram([binary, '--main_advice=MyMain', print_cmd])
    self.assertIn('MyMain', stdout)

    _, stdout, _ = self.RunProgram([
        binary,
        '--main_advice_classpath=foo/advice-1.jar;foo/advice-2.jar',
        print_cmd,
    ])
    self.assertIn('-classpath', stdout)
    classpath = stdout[stdout.index('-classpath') + 1]
    self.assertIn('foo/advice-1.jar', classpath)
    self.assertIn('foo/advice-2.jar', classpath)

    _, stdout, _ = self.RunProgram(
        [binary, '--main_advice_classpath=C:\\foo\\bar', print_cmd]
    )
    self.assertIn('-classpath', stdout)
    classpath = stdout[stdout.index('-classpath') + 1]
    self.assertIn('C:\\foo\\bar', classpath)

    _, stdout, _ = self.RunProgram(
        [binary, '--jvm_flag="--some_path="./a b/c""', print_cmd]
    )
    self.assertIn('"--some_path=\\"./a b/c\\""', stdout)

    _, stdout, _ = self.RunProgram(
        [binary, '--jvm_flags="--path1=a --path2=b"', print_cmd]
    )
    self.assertIn('--path1=a', stdout)
    self.assertIn('--path2=b', stdout)

    _, stdout, _ = self.RunProgram(
        [binary, print_cmd], env_add={'JVM_FLAGS': '--foo --bar'}
    )
    self.assertIn('--flag1', stdout)
    self.assertIn('--flag2', stdout)
    self.assertIn('--foo', stdout)
    self.assertIn('--bar', stdout)

    exit_code, stdout, stderr = self.RunProgram(
        [binary, '--singlejar', print_cmd], allow_failure=True
    )
    self.AssertExitCode(exit_code, 1, stderr)
    self.assertIn('foo_deploy.jar does not exist', ''.join(stderr))
    self.RunBazel(['build', '//foo:foo_deploy.jar'])
    _, stdout, _ = self.RunProgram([binary, '--singlejar', print_cmd])
    self.assertIn('-classpath', stdout)
    classpath = stdout[stdout.index('-classpath') + 1]
    self.assertIn('foo_deploy.jar', classpath)

    _, stdout, _ = self.RunProgram([binary, '--print_javabin'])
    self.assertIn('local_jdk/bin/java.exe', ''.join(stdout))

    my_tmp_dir = self.ScratchDir('my/temp/dir')
    _, stdout, _ = self.RunProgram(
        [binary, print_cmd], env_add={'TEST_TMPDIR': my_tmp_dir}
    )
    self.assertIn('-Djava.io.tmpdir=%s' % my_tmp_dir, stdout)

    _, stdout, _ = self.RunProgram([binary, '--classpath_limit=0', print_cmd])
    self.assertIn('-classpath', stdout)
    classpath = stdout[stdout.index('-classpath') + 1]
    self.assertRegexpMatches(classpath, r'foo-[A-Za-z0-9]+-classpath.jar$')

  def testWindowsNativeLauncherInNonEnglishPath(self):
    if not self.IsWindows():
      return
    self.CreateWorkspaceWithDefaultRepos('WORKSPACE')
    self.ScratchFile('bin/BUILD', [
        'java_binary(',
        '  name = "bin_java",',
        '  srcs = ["Main.java"],',
        '  main_class = "Main",',
        ')',
        'sh_binary(',
        '  name = "bin_sh",',
        '  srcs = ["main.sh"],',
        ')',
    ])
    self.ScratchFile('bin/Main.java', [
        'public class Main {',
        '  public static void main(String[] args) {'
        '    System.out.println("helloworld");',
        '  }',
        '}',
    ])
    self.ScratchFile('bin/main.sh', [
        'echo "helloworld"',
    ])

    _, stdout, _ = self.RunBazel(['info', 'bazel-bin'])
    bazel_bin = stdout[0]

    self.RunBazel(['build', '//bin/...'])

    for f in [
        'bin_java.exe',
        'bin_java.exe.runfiles_manifest',
        'bin_sh.exe',
        'bin_sh',
        'bin_sh.exe.runfiles_manifest',
    ]:
      self.CopyFile(os.path.join(bazel_bin, 'bin', f),
                    os.path.join(u'./\u6d4b\u8bd5', f))

    unicode_binary_path = u'./\u6d4b\u8bd5/bin_java.exe'
    _, stdout, _ = self.RunProgram([unicode_binary_path])
    self.assertEqual('helloworld', ''.join(stdout))

    unicode_binary_path = u'./\u6d4b\u8bd5/bin_sh.exe'
    _, stdout, _ = self.RunProgram([unicode_binary_path])
    self.assertEqual('helloworld', ''.join(stdout))

  def testWindowsNativeLauncherInLongPath(self):
    if not self.IsWindows():
      return
    self.CreateWorkspaceWithDefaultRepos('WORKSPACE')
    self.ScratchFile(
        'bin/BUILD',
        [
            'java_binary(',
            '  name = "not_short_bin_java",',
            '  srcs = ["Main.java"],',
            '  main_class = "Main",',
            ')',
            'sh_binary(',
            '  name = "not_short_bin_sh",',
            '  srcs = ["main.sh"],',
            ')',
            'py_binary(',
            '  name = "not_short_bin_py",',
            '  srcs = ["not_short_bin_py.py"],',
            ')',
        ],
    )
    self.ScratchFile('bin/Main.java', [
        'public class Main {',
        '  public static void main(String[] args) {'
        '    System.out.println("helloworld");',
        '  }',
        '}',
    ])
    self.ScratchFile('bin/main.sh', [
        'echo "helloworld"',
    ])
    self.ScratchFile(
        'bin/not_short_bin_py.py',
        [
            'print("helloworld")',
        ],
    )

    _, stdout, _ = self.RunBazel(['info', 'bazel-bin'])
    bazel_bin = stdout[0]

    exit_code, _, stderr = self.RunBazel(['build', '//bin/...'])
    self.AssertExitCode(exit_code, 0, stderr)

    # Create a directory with a path longer than 260
    long_dir_path = './' + '/'.join(
        [(c * 8 + '.' + c * 3) for c in string.ascii_lowercase])

    # The 'not_short_' prefix ensures that the basenames are not already 8.3
    # short paths. Due to the long directory path, the basename will thus be
    # replaced with a short path such as "not_sh~1.exe" below.
    for f in [
        'not_short_bin_java.exe',
        'not_short_bin_java.exe.runfiles_manifest',
        'not_short_bin_sh.exe',
        'not_short_bin_sh',
        'not_short_bin_sh.exe.runfiles_manifest',
        'not_short_bin_py.exe',
        'not_short_bin_py.zip',
        'not_short_bin_py.exe.runfiles_manifest',
    ]:
      self.CopyFile(
          os.path.join(bazel_bin, 'bin', f), os.path.join(long_dir_path, f))

    long_binary_path = os.path.abspath(
        long_dir_path + '/not_short_bin_java.exe'
    )
    # subprocess doesn't support long path without shell=True
    _, stdout, _ = self.RunProgram([long_binary_path], shell=True)
    self.assertEqual('helloworld', ''.join(stdout))
    # Make sure we can launch the binary with a shortened Windows 8dot3 path
    short_binary_path = win32api.GetShortPathName(long_binary_path)
    self.assertIn('~', os.path.basename(short_binary_path))
    _, stdout, _ = self.RunProgram([short_binary_path], shell=True)
    self.assertEqual('helloworld', ''.join(stdout))

    long_binary_path = os.path.abspath(long_dir_path + '/not_short_bin_sh.exe')
    # subprocess doesn't support long path without shell=True
    _, stdout, _ = self.RunProgram([long_binary_path], shell=True)
    self.assertEqual('helloworld', ''.join(stdout))
    # Make sure we can launch the binary with a shortened Windows 8dot3 path
    short_binary_path = win32api.GetShortPathName(long_binary_path)
    self.assertIn('~', os.path.basename(short_binary_path))
    _, stdout, _ = self.RunProgram([short_binary_path], shell=True)
    self.assertEqual('helloworld', ''.join(stdout))

    long_binary_path = os.path.abspath(long_dir_path + '/not_short_bin_py.exe')
    # subprocess doesn't support long path without shell=True
    _, stdout, _ = self.RunProgram([long_binary_path], shell=True)
    self.assertEqual('helloworld', ''.join(stdout))
    # Make sure we can launch the binary with a shortened Windows 8dot3 path
    short_binary_path = win32api.GetShortPathName(long_binary_path)
    self.assertIn('~', os.path.basename(short_binary_path))
    _, stdout, _ = self.RunProgram([short_binary_path], shell=True)
    self.assertEqual('helloworld', ''.join(stdout))

  def testWindowsNativeLauncherInvalidArgv0(self):
    if not self.IsWindows():
      return
    self.CreateWorkspaceWithDefaultRepos('WORKSPACE')
    self.ScratchFile(
        'bin/BUILD',
        [
            'java_binary(',
            '  name = "bin_java",',
            '  srcs = ["Main.java"],',
            '  main_class = "Main",',
            ')',
            'sh_binary(',
            '  name = "bin_sh",',
            '  srcs = ["main.sh"],',
            ')',
            'py_binary(',
            '  name = "bin_py",',
            '  srcs = ["bin_py.py"],',
            ')',
        ],
    )
    self.ScratchFile(
        'bin/Main.java',
        [
            'public class Main {',
            (
                '  public static void main(String[] args) {'
                '    System.out.println("helloworld");'
            ),
            '  }',
            '}',
        ],
    )
    self.ScratchFile(
        'bin/main.sh',
        [
            'echo "helloworld"',
        ],
    )
    self.ScratchFile(
        'bin/bin_py.py',
        [
            'print("helloworld")',
        ],
    )

    _, stdout, _ = self.RunBazel(['info', 'bazel-bin'])
    bazel_bin = stdout[0]

    self.RunBazel(['build', '//bin/...'])

    _, stdout, _ = self.RunProgram(
        ['C:\\Invalid'],
        executable=os.path.join(bazel_bin, 'bin', 'bin_java.exe'),
    )
    self.assertEqual('helloworld', ''.join(stdout))

    _, stdout, _ = self.RunProgram(
        ['C:\\Invalid'], executable=os.path.join(bazel_bin, 'bin', 'bin_sh.exe')
    )
    self.assertEqual('helloworld', ''.join(stdout))

    _, stdout, _ = self.RunProgram(
        ['C:\\Invalid'], executable=os.path.join(bazel_bin, 'bin', 'bin_py.exe')
    )
    self.assertEqual('helloworld', ''.join(stdout))

  def AssertRunfilesManifestContains(self, manifest, entry):
    with open(manifest, 'r') as f:
      for l in f:
        tokens = l.strip().split(' ', 1)
        if len(tokens) == 2 and tokens[0] == entry:
          return
    self.fail('Runfiles manifest "%s" did not contain "%s"' % (manifest, entry))

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


# -*- coding: utf-8 -*-

"""The Simulator takes in a :obj:`seagull.Board`, and runs a simulation given a
set number of iterations and a rule. For each iteration, the rule is applied to
the Board in order to evolve the lifeforms. After the simulation, run
statistics are returned.

.. code-block:: python

    import seagull as sg

    board = sg.Board()
    board.add(Blinker(), loc=(0,0))

    # Initialize a simulator
    sim = sg.Simulator(board)
    stats = sim.run(sg.rules.conway_classic, iters=1000)

You can always get the history of the whole simulation by calling the
`get_history()` method. The length of the history will always be equal to
:code:`iters + 1` since we include the initial state

.. note::

    Running a simulation does not change the :code:`state` attribute of the
    board. Internally, the simulator makes a copy of that layout and updates
    that instead. This is to avoid unintended behaviour when running
    simulations again and again.

Various statistics such as entropy, peak cell coverage, and the like are
returned as a dictionary. This gives us an idea on the characteristics of the
simulation experiment.

.. note::

    Some statistics are highly-dependent on the size of the board and the
    number of iterations. For example, peak cell coverage (pertaining to the
    max. amount of active cells during the whole run) depends on board size. If
    you have better ideas for computing these statistics, please open-up an
    Issue!

The :code:`run()` method only computes the progress of the board for the whole
simulation, but it does not animate it yet. To create an animation, call the
:code:`animate()` method:

.. code-block:: python

    sim.animate()

This returns a :obj:`matplotlib.animation.FuncAnimation` that you can turn into
an interactive animation in your notebook or exported as a GIF.

.. note::

    When exporting to GIF, it is required to have the ffmpeg backend installed.

"""

# Import standard library
from typing import Callable, Union

# Import modules
import matplotlib.pyplot as plt
import numpy as np
from loguru import logger
from matplotlib import animation

from .board import Board
from .utils import statistics as stats

class Simulator:
    def __init__(self, board: Board):
        """Initialize the class

        Parameters
        ----------
        board : seagull.Board
            The board to run the simulation on
        """
        self.board = board
        self.history = []  # type: list
        self.stats = {}  # type: dict

    def run(self, rule: Callable, iters: int, **kwargs) -> dict:
        """Run the simulation for a given number of iterations

        Parameters
        ----------
        rule : callable
            Callable that takes in an array and returns an array of the same
            shape.
        iters : int
            Number of iterations to run the simulation.

        Returns
        -------
        dict
           Computed statistics for the simulation run
        """
        layout = self.board.state.copy()

        # Append the initial state
        self.history.append(layout)

        # Run simulation
        for i in range(iters):
            layout = rule(layout, **kwargs)
            self.history.append(layout)

        self.stats = self.compute_statistics(self.get_history())
        return self.stats

    def compute_statistics(self, history: Union[list, np.ndarray]) -> dict:
        """Compute various statistics for the board

        Parameters
        ----------
        history : list or numpy.ndarray
            The simulation history

        Returns
        -------
        dict
            Compute statistics
        """
        logger.info("Computing simulation statistics...")
        sim_stats = {
            "peak_cell_coverage": np.max(
                [stats.cell_coverage(h) for h in history]
            ),
            "avg_cell_coverage": np.mean(
                [stats.cell_coverage(h) for h in history]
            ),
            "avg_shannon_entropy": np.mean(
                [stats.shannon_entropy(h) for h in history]
            ),
            "peak_shannon_entropy": np.max(
                [stats.shannon_entropy(h) for h in history]
            ),
        }

        return sim_stats

    def get_history(self, exclude_init=False) -> np.ndarray:
        """Get the simulation history

        Parameters
        ----------
        exclude_init: bool
            If True, then excludes the initial state in the history

        Returns
        -------
        numpy.ndarray
            Simulation history of shape :code:`(iters+1, board.size[0],
            board.size[1])`
        """
        history = self.history[1:] if exclude_init else self.history
        return np.asarray(history)

    def animate(self, figsize=(5, 5), interval=100) -> animation.FuncAnimation:
        """Animate the resulting simulation

        Parameters
        ----------
        figsize : tuple
            Size of the output figure
        interval : int
            Interval for transitioning between frames

        Returns
        -------
        matplotlib.animation.FuncAnimation
            Animation generated from the run
        """
        if not self.history:
            msg = "The run() argument must be executed first"
            logger.error(msg)
            raise ValueError(msg)

        logger.info("Rendering animation...")
        fig = plt.figure(figsize=figsize)
        ax = fig.add_axes([0, 0, 1, 1], xticks=[], yticks=[], frameon=False)
        X_blank = np.zeros(self.board.size, dtype=bool)
        im = ax.imshow(X_blank, cmap=plt.cm.binary, interpolation="nearest")
        im.set_clim(-0.05, 1)

        def _animate(i, history):
            current_pos = history[i]
            im.set_data(current_pos)
            return (im,)

        def _init():
            im.set_data(X_blank)
            return (im,)

        history = self.get_history()
        anim = animation.FuncAnimation(
            fig,
            func=_animate,
            frames=range(history.shape[0]),
            init_func=_init,
            interval=interval,
            fargs=(history,),
            blit=True,
        )
        return anim


import asyncio
from flask import Blueprint, abort, flash, redirect, render_template, request, jsonify, url_for, Response
from werkzeug.utils import secure_filename
import socket
from flask_socketio import SocketIO, emit
from app import app, db, socketio
import os
import time

HOST = "127.0.1.1"
WEBSOCKET_PORT = 9999
CHUNK_SIZE = 4096  # Define o tamanho do pacote. Pode ser ajustado conforme necessário.

# Lista de endereços IP dos servidores para armazenamento de réplicas
REPLICA_SERVERS = [HOST, HOST, HOST] #ips locais mockados
#REPLICA_SERVERS = ["192.168.1.2", "192.168.1.3", "192.168.1.4"]  # IPs das máquinas das réplicas

main = Blueprint('main', __name__)

MIME_TYPES = {
    "mp4": "video/mp4",
    "avi": "video/x-msvideo",
    "mkv": "video/x-matroska",
    "flv": "video/x-flv"
}

class StreamingError(Exception):
    """Exceção personalizada para erros de streaming."""
    pass

class Video(db.Model):
    __tablename__ = 'video'
    __table_args__ = {'extend_existing': True}

    id = db.Column(db.Integer, primary_key=True)
    filename = db.Column(db.String(150), unique=True, nullable=False)
    description = db.Column(db.String(500), nullable=True)

with app.app_context():
    db.create_all()

def allowed_file(filename):
    return '.' in filename and filename.rsplit('.', 1)[1].lower() in app.config['ALLOWED_EXTENSIONS']

def upload_to_replica(filename, file_content):
    for server_ip in REPLICA_SERVERS:
        try:
            client = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
            client.connect((server_ip, WEBSOCKET_PORT))
            # Enviar comando UPLOAD
            header = f"UPLOAD"
            client.send(header.encode())
            # Enviar tamanho do arquivo como uma string de tamanho 10
            client.send(str(len(file_content)).encode().zfill(10))
            # Enviar tamanho do nome do arquivo
            client.send(str(len(filename)).encode().zfill(10))
            # Enviar nome do arquivo
            client.send(filename.encode())
            # Enviar os dados do arquivo
            client.sendall(file_content)
            client.close()
        except Exception as e:
            print(f"Erro ao enviar para servidor {server_ip}: {e}")

@main.route('/upload', methods=['POST'])
def upload_file():
    if 'file' not in request.files:
        return jsonify({"error": "No file provided"}), 400
    file = request.files['file']
    if file.filename == '':
        return jsonify({"error": "No file selected"}), 400
    if file and allowed_file(file.filename):
        filename = secure_filename(file.filename)
        file_path = os.path.join(app.config['UPLOAD_FOLDER'], filename)
        file.save(file_path)
        new_video = Video(filename=filename)
        db.session.add(new_video)
        db.session.commit()

        with open(file_path, 'rb') as f:
            file_content = f.read()
        upload_to_replica(filename, file_content)
        return "File uploaded successfully! You can now upload another file."
    return jsonify({"error": "Invalid file type"}), 400

@main.route('/', methods=['GET'])
def show_upload():
    return render_template('upload.html')

@main.route('/videos', methods=['GET'])
def list_videos():
    videos = Video.query.all()
    return render_template('video_list.html', videos=videos)

from websockets import connect as ws_connect

@main.route('/play/<int:video_id>', methods=['GET'])
def play_video(video_id):
    video = Video.query.get(video_id)
    video_name = video.filename

    # Adicionando failover para o streaming de vídeo
    for _ in range(3):  # Tenta até 3 vezes, uma para cada réplica
        try:
            return stream_video(video_name)
        except StreamingError:
            continue  # Se ocorrer um erro, tenta a próxima réplica

    return "Não foi possível reproduzir o vídeo."

def stream_video(video_name):
    try:
        client = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        client.connect((HOST, 9999))

        header = f"STREAM"
        client.send(header.encode())
        client.send(str(len(video_name)).zfill(10).encode())
        client.send(video_name.encode())

        def generate():
            while True:
                chunk = client.recv(CHUNK_SIZE)
                if not chunk:
                    break
                yield chunk

        ext = video_name.split('.')[-1]
        mime_type = MIME_TYPES.get(ext, "video/mp4")
        return Response(generate(), content_type=mime_type)

    except ConnectionError:
        # Esta exceção pode ser lançada se houver um problema de conexão de rede
        raise StreamingError("Erro de conexão durante o streaming do vídeo")

@main.route('/delete_video/<int:video_id>', methods=['POST'])
def delete_video(video_id):
    video = Video.query.get(video_id)
    if video:
        db.session.delete(video)
        db.session.commit()
        return redirect(url_for('main.list_videos'))
    else:
        # Caso o vídeo não seja encontrado no banco de dados
        flash('Vídeo não encontrado', 'error')
        return redirect(url_for('main.list_videos'))

if __name__ == '__main__':
    app.run(debug=True)


"""
3D convolutions using GPU accelereration for Theano (using conv2d)
https://github.com/jaberg/TheanoConv3d2d

"""

import theano
from theano.gradient import DisconnectedType
from theano.gof import Op, Apply
from theano import tensor
import theano.sandbox.cuda as cuda

def get_diagonal_subtensor_view(x, i0, i1):
    """Helper function for DiagonalSubtensor and
    IncDiagonalSubtensor

    :note: it return a partial view of x, not a partial copy.
    """

    if x.shape[i0] < x.shape[i1]:
        raise NotImplementedError('is this allowed?')
    idx = [slice(None)] * x.ndim
    idx[i0] = slice(x.shape[i1] - 1, None, None)
    xview = x.__getitem__(tuple(idx))
    strides = list(xview.strides)
    strides[i1] -= strides[i0]
    xview.strides = strides
    return xview

class DiagonalSubtensor(Op):
    """Return a form a nd diagonal subtensor.

    :param x: n-d tensor
    :param i0: axis index in x
    :param i1: axis index in x
    :note: Work on the GPU.

    ``x`` is some n-dimensional tensor, but this Op only deals with a
    matrix-shaped slice, using axes i0 and i1. Without loss of
    generality, suppose that ``i0`` picks out our ``row`` dimension,
    and i1 the ``column`` dimension.

    So the relevant part of ``x`` is some matrix ``u``. Suppose it has 7 rows
    and 4 columns::

        [ 0 0 0 0 ]
        [ 0 0 0 0 ]
        [ 0 0 0 0 ]
        [ 0 0 0 0 ]
        [ 0 0 0 0 ]
        [ 0 0 0 0 ]

    The view returned by this function is also a matrix. It's a thick,
    diagonal ``stripe`` across u that discards the lower left triangle
    and the upper right triangle:

        [ x 0 0 0 ]
        [ x x 0 0 ]
        [ x x x 0 ]
        [ 0 x x x ]
        [ 0 0 x x ]
        [ 0 0 0 x ]

    In this case the return value would be this view of shape 3x4. The
    returned view has the same number of dimensions as the input
    ``x``, and the only difference is that the shape along dimension
    ``i0`` has been reduced by ``shape[i1] - 1`` because of the
    triangles that got chopped out.

    The NotImplementedError is meant to catch the case where shape[i0]
    is too small for the stripe to reach across the matrix, in which
    case it's not clear what this function should do. Maybe always
    raise an error. I'd look back to the call site in the Conv3D to
    see what's necessary at that point.

    """
    def __str__(self):
        if self.inplace:
            return "%s{inplace}" % self.__class__.__name__
        return "%s" % self.__class__.__name__

    def __init__(self, inplace=False):
        self.inplace = inplace
        if inplace:
            self.view_map = {0: [0]}

    def __eq__(self, other):
        return type(self) == type(other) and self.inplace == other.inplace

    def __hash__(self):
        return hash((type(self), self.inplace))

    def make_node(self, x, i0, i1):
        _i0 = tensor.as_tensor_variable(i0)
        _i1 = tensor.as_tensor_variable(i1)
        return Apply(self, [x, _i0, _i1], [x.type()])

    def perform(self, node, inputs, output_storage):
        xview = get_diagonal_subtensor_view(*inputs)
        if self.inplace:
            output_storage[0][0] = xview
        else:
            output_storage[0][0] = xview.copy()

    def grad(self, inputs, g_outputs):
        z = tensor.zeros_like(inputs[0])
        gx = inc_diagonal_subtensor(z, inputs[1], inputs[2], g_outputs[0])
        return [gx, DisconnectedType()(), DisconnectedType()()]

    def connection_pattern(self, node):
        rval = [[True], [False], [False]]
        return rval

diagonal_subtensor = DiagonalSubtensor(False)

class IncDiagonalSubtensor(Op):
    """
    The gradient of DiagonalSubtensor
    """
    def __str__(self):
        if self.inplace:
            return "%s{inplace}" % self.__class__.__name__
        return "%s" % self.__class__.__name__

    def __init__(self, inplace=False):
        self.inplace = inplace
        if inplace:
            self.destroy_map = {0: [0]}

    def __eq__(self, other):
        return type(self) == type(other) and self.inplace == other.inplace

    def __hash__(self):
        return hash((type(self), self.inplace))

    def make_node(self, x, i0, i1, amt):
        _i0 = tensor.as_tensor_variable(i0)
        _i1 = tensor.as_tensor_variable(i1)
        return Apply(self, [x, _i0, _i1, amt], [x.type()])

    def perform(self, node, inputs, output_storage):
        x, i0, i1, amt = inputs
        if not self.inplace:
            x = x.copy()
        xview = get_diagonal_subtensor_view(x, i0, i1)
        xview += amt
        output_storage[0][0] = x

    def grad(self, inputs, g_outputs):
        x, i0, i1, amt = inputs
        gy = g_outputs[0]
        return [gy, DisconnectedType()(), DisconnectedType()(),
                diagonal_subtensor(gy, i0, i1)]

    def connection_pattern(self, node):
        rval = [[True], [False], [False], [True]]
        return rval
inc_diagonal_subtensor = IncDiagonalSubtensor(False)

def conv3d(signals, filters,
           signals_shape=None, filters_shape=None,
           border_mode='valid'):
    """Convolve spatio-temporal filters with a movie.

    :param signals: timeseries of images whose pixels have color channels.
            shape: [Ns, Ts, C, Hs, Ws]
    :param filters: spatio-temporal filters
            shape: [Nf, Tf, C, Hf, Wf]
    :param signals_shape: None or a tuple/list with the shape of signals
    :param filters_shape: None or a tuple/list with the shape of filters
    :param border_mode: The only one tested is 'valid'.

    :note: Work on the GPU.
    """

    if isinstance(border_mode, str):
        border_mode = (border_mode, border_mode, border_mode)

    _signals_shape_5d = signals.shape if signals_shape is None else signals_shape
    _filters_shape_5d = filters.shape if filters_shape is None else filters_shape

    _signals_shape_4d = (
        _signals_shape_5d[0] * _signals_shape_5d[1],
        _signals_shape_5d[2],
        _signals_shape_5d[3],
        _signals_shape_5d[4],
        )
    _filters_shape_4d = (
        _filters_shape_5d[0] * _filters_shape_5d[1],
        _filters_shape_5d[2],
        _filters_shape_5d[3],
        _filters_shape_5d[4],
        )

    if border_mode[1] != border_mode[2]:
        raise NotImplementedError('height and width bordermodes must match')
    conv2d_signal_shape = _signals_shape_4d
    conv2d_filter_shape = _filters_shape_4d
    if signals_shape is None:
        conv2d_signal_shape = None
    if filters_shape is None:
        conv2d_filter_shape = None

    out_4d = tensor.nnet.conv2d(
        signals.reshape(_signals_shape_4d),
        filters.reshape(_filters_shape_4d),
        image_shape=conv2d_signal_shape,
        filter_shape=conv2d_filter_shape,
        border_mode = border_mode[1])  # ignoring border_mode[2]

    # reshape the output to restore its original size
    # shape = Ns, Ts, Nf, Tf, W-Wf+1, H-Hf+1
    if border_mode[1] == 'valid':
        out_tmp = out_4d.reshape((
            _signals_shape_5d[0],  # Ns
            _signals_shape_5d[1],  # Ts
            _filters_shape_5d[0],  # Nf
            _filters_shape_5d[1],  # Tf
            _signals_shape_5d[3] - _filters_shape_5d[3] + 1,
            _signals_shape_5d[4] - _filters_shape_5d[4] + 1,
            ))
    elif border_mode[1] == 'full':
        out_tmp = out_4d.reshape((
            _signals_shape_5d[0],  # Ns
            _signals_shape_5d[1],  # Ts
            _filters_shape_5d[0],  # Nf
            _filters_shape_5d[1],  # Tf
            _signals_shape_5d[3] + _filters_shape_5d[3] - 1,
            _signals_shape_5d[4] + _filters_shape_5d[4] - 1,
            ))
    elif border_mode[1] == 'same':
        raise NotImplementedError()
    else:
        raise ValueError('invalid border mode', border_mode[1])

    # now sum out along the Tf to get the output
    # but we have to sum on a diagonal through the Tf and Ts submatrix.
    if border_mode[0] == 'valid':
        out_5d = diagonal_subtensor(out_tmp, 1, 3).sum(axis=3)
    elif border_mode[0] in ('full', 'same'):
        out_5d = out_4d.reshape((_signals_shape_5d))
        # raise NotImplementedError('sequence border mode', border_mode[0])
    else:
        raise ValueError('invalid border mode', border_mode[1])
    return out_5d

def make_gpu_optimizer(op, to_gpu):
    """This function create optimizer that move some inputs to the GPU
    for op that work on both CPU and GPU.

    The op object is created by calling op(), so good default value
    are needed.

    We suppose the same op work with CPU and GPU inputs.

    :param op: the op that support GPU inputs
    :param to_gpu: a list of op inputs that are moved to the GPU.

    """
    @theano.gof.local_optimizer([])
    def local_to_gpu(node):
        """
        op(host_from_gpu()) -> host_from_gpu(op)
        gpu_from_host(op) -> op(gpu_from_host)
        """
        if isinstance(node.op, op):
            #op(host_from_gpu()) -> host_from_gpu(op)
            #If any of the input that go on the GPU are on the GPU,
            #move the op to the gpu.
            if any(node.inputs[idx].owner and
                   isinstance(node.inputs[idx].owner.op, cuda.HostFromGpu)
                   for idx in to_gpu):
                new_inp = list(node.inputs)
                for idx in to_gpu:
                    new_inp[idx] = cuda.gpu_from_host(new_inp[idx])
                return [cuda.host_from_gpu(op()(*new_inp))]
        if node.op == cuda.gpu_from_host:
            #gpu_from_host(op) -> op(gpu_from_host)
            host_input = node.inputs[0]
            if host_input.owner and isinstance(host_input.owner.op,
                                               op):
                op_node = host_input.owner
                new_inp = list(op_node.inputs)
                for idx in to_gpu:
                    new_inp[idx] = cuda.gpu_from_host(new_inp[idx])
                return [op()(*new_inp)]
        return False
    local_to_gpu.__name__ = "local_to_gpu_" + op.__name__
    cuda.opt.register_opt()(local_to_gpu)

if cuda.cuda_available:
    make_gpu_optimizer(DiagonalSubtensor, [0])
    make_gpu_optimizer(IncDiagonalSubtensor, [0, 3])


import math
import copy
import numpy as np
import pprint
import torch
import torch.nn as nn
import torch.nn.functional as F
from fvcore.nn.precise_bn import get_bn_modules, update_bn_stats

import slowfast.models.losses as losses
import slowfast.models.optimizer as optim
import slowfast.utils.checkpoint as cu
import slowfast.utils.distributed as du
import slowfast.utils.logging as logging
import slowfast.utils.metrics as metrics
import slowfast.utils.misc as misc
import slowfast.visualization.tensorboard_vis as tb
from slowfast.datasets import loader
from slowfast.datasets.mixup import MixUp
from slowfast.models import build_model
from slowfast.utils.meters import EpochTimer, TrainMeter, ValMeter, AdaMeter

logger = logging.get_logger(__name__)

def train_epoch(
    train_loaders,
    model,
    optimizers,
    scaler,
    train_meter,
    cur_epoch,
    cfg,
    writer=None,
):
    """
    Perform the video training for one epoch.
    Args:
        train_loaders (list of loader): source and target video training loader.
        model (model): the video model to train.
        optimizer (optim): the optimizer to perform optimization on the model's
            parameters.
        train_meter (TrainMeter): training meters to log the training performance.
        cur_epoch (int): current epoch of training.
        cfg (CfgNode): configs. Details can be found in
            slowfast/config/defaults.py
        writer (TensorboardWriter, optional): TensorboardWriter object
            to writer Tensorboard log.
    """
    source_loader = train_loaders[0]
    target_unl_loader = train_loaders[1]
    if cfg.ADAPTATION.SEMI_SUPERVISED.ENABLE:
        target_lab_loader = train_loaders[2]
    optimizer_f, optimizer_c = optimizers[0], optimizers[1]

    # Enable train mode.
    model.train()
    train_meter.iter_tic()
    data_size = len(source_loader)
    target_unl_iter = iter(target_unl_loader)
    target_unl_size = len(target_unl_loader)
    if cfg.ADAPTATION.SEMI_SUPERVISED.ENABLE:
        target_lab_iter = iter(target_lab_loader)
        target_lab_size = len(target_lab_loader)

    for cur_iter, (inputs_source, labels_source, _, _) in enumerate(source_loader):
        # Load the data.
        if cur_iter%target_unl_size==0:
            target_unl_iter = iter(target_unl_loader)
        inputs_target_unl, labels_target_unl, _, _ = next(target_unl_iter)
        if cfg.ADAPTATION.SEMI_SUPERVISED.ENABLE:
            if cur_iter%target_lab_size==0:
                target_lab_iter = iter(target_lab_loader)
            inputs_target_lab, labels_target_lab, _, _ = next(target_lab_iter)

        # Transfer the data to the current GPU device.
        for i in range(len(inputs_source)):
            inputs_source[i] = inputs_source[i].cuda(non_blocking=True)
            inputs_target_unl[i] = inputs_target_unl[i].cuda(non_blocking=True)
        labels_source = labels_source.cuda()
        labels_target_unl = labels_target_unl.cuda()
        if cfg.ADAPTATION.SEMI_SUPERVISED.ENABLE:
            for i in range(len(inputs_source)):
                inputs_target_lab[i] = inputs_target_lab[i].cuda(non_blocking=True)
            labels_target_lab = labels_target_lab.cuda()

        # Update the learning rate.
        lr = optim.get_epoch_lr(cur_epoch + float(cur_iter) / data_size, cfg)
        optim.set_lr(optimizer_f, lr)
        optim.set_lr(optimizer_c, lr)

        train_meter.data_toc()
        source_weak = inputs_source[1]
        source_strong = inputs_source[0]
        target_unl_weak = inputs_target_unl[1]
        target_unl_strong = inputs_target_unl[0]
        if cfg.ADAPTATION.SEMI_SUPERVISED.ENABLE:
            target_lab_weak = inputs_target_lab[1]
            target_lab_strong = inputs_target_lab[0]

        if not cfg.ADAPTATION.SEMI_SUPERVISED.ENABLE:
            lab_inputs = [source_strong]
            lab_labels = labels_source
            unl_inputs = [target_unl_weak]
            unl_labels = labels_target_unl
        else:
            lab_inputs = [torch.cat((source_strong, target_lab_strong), dim=0)]
            lab_labels = torch.cat((labels_source, labels_target_lab), dim=0)
            unl_inputs = [target_unl_weak]
            unl_labels = labels_target_unl

        with torch.cuda.amp.autocast(enabled=cfg.TRAIN.MIXED_PRECISION):
            # Step A train all networks to minimize loss on source domain
            optimizer_f.zero_grad()
            optimizer_c.zero_grad()

            lab_preds, lab_feats = model(lab_inputs)

            criterion = nn.CrossEntropyLoss()
            loss_s = criterion(lab_preds, lab_labels)
            loss_s.backward()
            optimizer_f.step()
            optimizer_c.step()

            # Step B train classifier to maximize discrepancy
            optimizer_f.zero_grad()
            optimizer_c.zero_grad()

            unl_preds, unl_feats = model(unl_inputs, reverse=True)
            new_preds = F.softmax(unl_preds, dim=1)
            loss_h = cfg.MME.LAMBDA * torch.mean(
                torch.sum(new_preds * (torch.log(new_preds + 1e-5)), 1))
            loss_h.backward()
            optimizer_f.step()
            optimizer_c.step()

            prototypes = model.module.head.weight.clone().detach()

        # Compute the errors.
        num_topks_correct = metrics.topks_correct(lab_preds, lab_labels, (1, 5))
        top1_err, top5_err = [
            (1.0 - x / lab_preds.size(0)) * 100.0 for x in num_topks_correct
        ]

        # Gather all the predictions across all the devices.
        if cfg.NUM_GPUS > 1:
            loss_s, loss_h, top1_err, top5_err = du.all_reduce(
                [loss_s, loss_h, top1_err, top5_err]
            )

        # Copy the stats from GPU to CPU (sync point).
        loss_s, loss_h, top1_err, top5_err = (
            loss_s.item(),
            loss_h.item(),
            top1_err.item(),
            top5_err.item()
        )
        batch_size = inputs_source[0].size(0)*max(cfg.NUM_GPUS, 1)

        # Update and log stats.
        train_meter.update_stats(
            top1_err,
            top5_err,
            loss_s,
            lr,
            batch_size,
        )
        # write to tensorboard format if available.
        if writer is not None:
            dict2write = {
                "Train/loss_s": loss_s,
                "Train/loss_h": -loss_h,
                "Train/lr": lr,
                "Train/Top1_err": top1_err,
                "Train/Top5_err": top5_err,
            }
            writer.add_scalars(dict2write, global_step=data_size * cur_epoch + cur_iter)

            if cfg.TENSORBOARD.DIST_VIS.ENABLE and (data_size * cur_epoch + cur_iter)%cfg.TENSORBOARD.DIST_VIS.LOG_PERIOD==1:
                writer.add_confusion_matrix(
                    torch.argmax(torch.cat(train_meter.all_source_weak, dim=0), dim=1),
                    torch.cat(train_meter.all_source_labels, dim=0),
                    tag="Confusion/Labeled",
                    global_step=data_size * cur_epoch + cur_iter
                )
                writer.add_confusion_matrix(
                    torch.argmax(torch.cat(train_meter.all_target_weak, dim=0), dim=1),
                    torch.cat(train_meter.all_target_labels, dim=0),
                    tag="Confusion/Unlabeled",
                    global_step=data_size * cur_epoch + cur_iter
                )

            if cfg.TENSORBOARD.SAMPLE_VIS.ENABLE and (data_size * cur_epoch + cur_iter)%cfg.TENSORBOARD.SAMPLE_VIS.LOG_PERIOD==0:
                writer.add_video_pred(
                    lab_inputs[0],
                    torch.argmax(lab_preds, dim=1),
                    lab_labels,
                    tag="Sample/Source",
                    global_step = data_size * cur_epoch + cur_iter,
                )
                writer.add_video_pred(
                    unl_inputs[0],
                    torch.argmax(unl_preds, dim=1),
                    unl_labels,
                    tag="Sample/Target",
                    global_step = data_size * cur_epoch + cur_iter,
                )

        train_meter.iter_toc()  # measure allreduce for this meter
        train_meter.update_predictions(
            lab_preds, lab_feats, lab_labels,
            unl_preds, unl_feats, unl_labels, prototypes,
        )
        train_meter.log_iter_stats(cur_epoch, cur_iter)
        torch.cuda.synchronize()
        train_meter.iter_tic()
        del inputs_source, inputs_target_unl, labels_source, labels_target_unl
        if cfg.ADAPTATION.SEMI_SUPERVISED.ENABLE:
            del inputs_target_lab, labels_target_lab

        # in case of fragmented memory
        torch.cuda.empty_cache()

    # Log epoch stats.
    train_meter.log_epoch_stats(cur_epoch)
    # write to tensorboard format if available.
    if writer is not None:
        if cfg.TENSORBOARD.EPOCH_LOG.ENABLE:
            writer.writer.add_scalars(
                "Error/Top1_err",
                {"Train": train_meter.num_top1_mis / train_meter.num_samples}, global_step=cur_epoch
            )
            writer.writer.add_scalars(
                "Error/Top5_err",
                {"Train": train_meter.num_top5_mis / train_meter.num_samples}, global_step=cur_epoch
            )
        if cfg.TENSORBOARD.CONFUSION_MATRIX.ENABLE:
            all_preds = [pred.clone().detach() for pred in train_meter.all_source_strong]
            all_labels = [label.clone().detach() for label in train_meter.all_source_labels]
            all_preds = [pred.cpu() for pred in all_preds]
            all_labels = [label.cpu() for label in all_labels]
            writer.plot_eval(
                preds=all_preds,
                labels=all_labels,
                global_step=cur_epoch,
                tag="Confusion/Train"
            )
    train_meter.reset()

@torch.no_grad()
def eval_epoch(
    val_loader, model, val_meter, cur_epoch, cfg, writer=None
    ):
    """
    Evaluate the model on the val set.
    Args:
        val_loader (loader): data loader to provide validation data.
        model (model): model to evaluate the performance.
        val_meter (ValMeter): meter instance to record and calculate the metrics.
        cur_epoch (int): number of the current epoch of training.
        cfg (CfgNode): configs. Details can be found in
            slowfast/config/defaults.py
        writer (TensorboardWriter, optional): TensorboardWriter object
            to writer Tensorboard log.
    """

    # Evaluation mode enabled. The running stats would not be updated.
    model.eval()
    val_meter.iter_tic()

    for cur_iter, (inputs, labels, _, meta) in enumerate(val_loader):
        if cfg.NUM_GPUS:
            # Transferthe data to the current GPU device.
            if isinstance(inputs, (list,)):
                for i in range(len(inputs)):
                    inputs[i] = inputs[i].cuda(non_blocking=True)
            else:
                inputs = inputs.cuda(non_blocking=True)
            labels = labels.cuda()
            for key, val in meta.items():
                if isinstance(val, (list,)):
                    for i in range(len(val)):
                        val[i] = val[i].cuda(non_blocking=True)
                else:
                    meta[key] = val.cuda(non_blocking=True)
        val_meter.data_toc()

        preds, _ = model(inputs)

        if cfg.DATA.MULTI_LABEL:
            if cfg.NUM_GPUS > 1:
                preds, labels = du.all_gather([preds, labels])
        else:
            # Compute the errors.
            num_topks_correct = metrics.topks_correct(preds, labels, (1, 5))

            # Combine the errors across the GPUs.
            top1_err, top5_err = [
                (1.0 - x / preds.size(0)) * 100.0 for x in num_topks_correct
            ]
            if cfg.NUM_GPUS > 1:
                top1_err, top5_err = du.all_reduce([top1_err, top5_err])

            # Copy the errors from GPU to CPU (sync point).
            top1_err, top5_err = top1_err.item(), top5_err.item()

            val_meter.iter_toc()
            # Update and log stats.
            val_meter.update_stats(
                top1_err,
                top5_err,
                inputs[0].size(0)
                * max(
                    cfg.NUM_GPUS, 1
                ),  # If running  on CPU (cfg.NUM_GPUS == 1), use 1 to represent 1 CPU.
            )
            # write to tensorboard format if available.
            if writer is not None:
                writer.add_scalars(
                    {"Val/Top1_err": top1_err, "Val/Top5_err": top5_err},
                    global_step=len(val_loader) * cur_epoch + cur_iter,
                )

                if cfg.TENSORBOARD.SAMPLE_VIS.ENABLE and (len(val_loader) * cur_epoch + cur_iter)%cfg.TENSORBOARD.SAMPLE_VIS.LOG_PERIOD==0:
                    writer.add_video_pred(
                        inputs[0],
                        torch.argmax(preds, dim=1),
                        labels,
                        tag="Sample/Val",
                        global_step = len(val_loader) * cur_epoch + cur_iter,
                    )

            val_meter.update_predictions(preds, labels)

        val_meter.log_iter_stats(cur_epoch, cur_iter)
        val_meter.iter_tic()

    # Log epoch stats.
    val_meter.log_epoch_stats(cur_epoch)
    # write to tensorboard format if available.
    if writer is not None:
        if cfg.TENSORBOARD.EPOCH_LOG.ENABLE:
            writer.writer.add_scalars(
                "Error/Top1_err",
                {"Val": val_meter.num_top1_mis / val_meter.num_samples}, global_step=cur_epoch
            )
            writer.writer.add_scalars(
                "Error/Top5_err",
                {"Val": val_meter.num_top5_mis / val_meter.num_samples}, global_step=cur_epoch
            )
        if cfg.TENSORBOARD.CONFUSION_MATRIX.ENABLE:
            all_preds = [pred.clone().detach() for pred in val_meter.all_preds]
            all_labels = [
                label.clone().detach() for label in val_meter.all_labels
            ]
            if cfg.NUM_GPUS:
                all_preds = [pred.cpu() for pred in all_preds]
                all_labels = [label.cpu() for label in all_labels]
            writer.plot_eval(
                preds=all_preds,
                labels=all_labels,
                global_step=cur_epoch,
                tag="Confusion/Val"
            )

    val_meter.reset()

def calculate_and_update_precise_bn(loader, model, num_iters=200, use_gpu=True):
    """
    Update the stats in bn layers by calculate the precise stats.
    Args:
        loader (loader): data loader to provide training data.
        model (model): model to update the bn stats.
        num_iters (int): number of iterations to compute and update the bn stats.
        use_gpu (bool): whether to use GPU or not.
    """

    def _gen_loader():
        for inputs, *_ in loader:
            if use_gpu:
                if isinstance(inputs, (list,)):
                    for i in range(len(inputs)):
                        inputs[i] = inputs[i].cuda(non_blocking=True)
                else:
                    inputs = inputs.cuda(non_blocking=True)
            yield inputs

    # Update the bn stats.
    update_bn_stats(model, _gen_loader(), num_iters)

def train(cfg):
    """
    Train a video model for many epochs on train set and evaluate it on val set.
    Args:
        cfg (CfgNode): configs. Details can be found in
            slowfast/config/defaults.py
    """
    # Set up environment.
    du.init_distributed_training(cfg)
    # Set random seed from configs.
    np.random.seed(cfg.RNG_SEED)
    torch.manual_seed(cfg.RNG_SEED)

    # Setup logging format.
    logging.setup_logging(cfg.OUTPUT_DIR)

    # Print config.
    logger.info("Train with config:")
    logger.info(pprint.pformat(cfg))

    # Build the video model and print model statistics.
    cfg.EXTRACT.ENABLE = True
    cfg.SWIN.TEMP = cfg.MME.TEMP
    cfg.SWIN.ETA = cfg.MME.ETA
    model = build_model(cfg)
    if du.is_master_proc() and cfg.LOG_MODEL_INFO:
        misc.log_model_info(model, cfg, use_train_input=True)

    # Construct the optimizer.
    sub_modules = []
    if cfg.NUM_GPUS > 1:
        for name, sub_module in model.module.named_modules():
            if name!="head":
                sub_modules.append(sub_module)
    else:
        for name, sub_module in model.named_modules():
            if name!="head":
                sub_modules.append(sub_module)
    backbone = nn.Sequential(*sub_modules)
    classifier = model.module.get_submodule("head")
    optimizer_f = optim.construct_optimizer(backbone, cfg)
    optimizer_c = optim.construct_optimizer(classifier, cfg)
    optimizers = [optimizer_f, optimizer_c]
    # Create a GradScaler for mixed precision training
    scaler = torch.cuda.amp.GradScaler(enabled=cfg.TRAIN.MIXED_PRECISION)

    # Load a checkpoint to resume training if applicable.
    start_epoch = cu.load_train_checkpoint(cfg, model, optimizer_f,
        scaler if cfg.TRAIN.MIXED_PRECISION else None)

    # Create the video train and val loaders.
    if cfg.ADAPTATION.SEMI_SUPERVISED.ENABLE:
        source_cfg = copy.deepcopy(cfg)
        source_cfg.DATA.IMDB_FILES.TRAIN = cfg.ADAPTATION.SOURCE
        source_cfg.DATA.IMDB_FILES.VAL = cfg.ADAPTATION.TARGET
        source_loader = loader.construct_loader(source_cfg, "train")
        val_loader = loader.construct_loader(source_cfg, "val")
        target_lab_cfg = copy.deepcopy(cfg)
        target_lab_cfg.DATA.IMDB_FILES.TRAIN = cfg.ADAPTATION.TARGET
        target_lab_cfg.DATA.IMDB_FILES.VAL = cfg.ADAPTATION.SOURCE
        target_lab_cfg.TRAIN.BATCH_SIZE = int(cfg.ADAPTATION.ALPHA * source_cfg.TRAIN.BATCH_SIZE)
        target_lab_loader = loader.construct_loader(target_lab_cfg, "lab")
        target_unl_cfg = copy.deepcopy(cfg)
        target_unl_cfg.DATA.IMDB_FILES.TRAIN = cfg.ADAPTATION.TARGET
        target_unl_cfg.DATA.IMDB_FILES.VAL = cfg.ADAPTATION.SOURCE
        target_unl_cfg.TRAIN.BATCH_SIZE = int(cfg.ADAPTATION.BETA * source_cfg.TRAIN.BATCH_SIZE)
        target_unl_loader = loader.construct_loader(target_unl_cfg, "unl")
        bn_cfg = copy.deepcopy(cfg)
        bn_cfg.DATA.IMDB_FILES.TRAIN = cfg.ADAPTATION.SOURCE + cfg.ADAPTATION.TARGET
        bn_cfg.ADAMATCH.ENABLE = False
        precise_bn_loader = (
            loader.construct_loader(bn_cfg, "train", is_precise_bn=True)
            if cfg.BN.USE_PRECISE_STATS
            else None
        )
        train_loaders = [source_loader, target_unl_loader, target_lab_loader]
    else:
        source_cfg = copy.deepcopy(cfg)
        source_cfg.DATA.IMDB_FILES.TRAIN = cfg.ADAPTATION.SOURCE
        source_cfg.DATA.IMDB_FILES.VAL = cfg.ADAPTATION.TARGET
        source_loader = loader.construct_loader(source_cfg, "train")
        val_loader = loader.construct_loader(source_cfg, "val")
        target_unl_cfg = copy.deepcopy(cfg)
        target_unl_cfg.DATA.IMDB_FILES.TRAIN = cfg.ADAPTATION.TARGET
        target_unl_cfg.DATA.IMDB_FILES.VAL = cfg.ADAPTATION.SOURCE
        target_unl_cfg.TRAIN.BATCH_SIZE = int(cfg.ADAPTATION.BETA * source_cfg.TRAIN.BATCH_SIZE)
        target_unl_loader = loader.construct_loader(target_unl_cfg, "train")
        bn_cfg = copy.deepcopy(cfg)
        bn_cfg.DATA.IMDB_FILES.TRAIN = cfg.ADAPTATION.SOURCE + cfg.ADAPTATION.TARGET
        bn_cfg.ADAMATCH.ENABLE = False
        precise_bn_loader = (
            loader.construct_loader(bn_cfg, "train", is_precise_bn=True)
            if cfg.BN.USE_PRECISE_STATS
            else None
        )
        train_loaders = [source_loader, target_unl_loader]

    # Create meters.
    train_meter = AdaMeter(len(train_loaders[0]), cfg)
    val_meter = ValMeter(len(val_loader), cfg)

    # set up writer for logging to Tensorboard format.
    if cfg.TENSORBOARD.ENABLE and du.is_master_proc(
        cfg.NUM_GPUS * cfg.NUM_SHARDS
    ):
        writer = tb.TensorboardWriter(cfg)
    else:
        writer = None

    # Perform the training loop.
    logger.info("Start epoch: {}".format(start_epoch + 1))

    epoch_timer = EpochTimer()
    for cur_epoch in range(start_epoch, cfg.SOLVER.MAX_EPOCH):

        # Shuffle the dataset.
        for train_loader in train_loaders:
            loader.shuffle_dataset(train_loader, cur_epoch)
        # Train for one epoch.
        epoch_timer.epoch_tic()
        train_epoch(
            train_loaders,
            model,
            optimizers,
            scaler,
            train_meter,
            cur_epoch,
            cfg,
            writer,
        )
        epoch_timer.epoch_toc()
        logger.info(
            f"Epoch {cur_epoch} takes {epoch_timer.last_epoch_time():.2f}s. Epochs "
            f"from {start_epoch} to {cur_epoch} take "
            f"{epoch_timer.avg_epoch_time():.2f}s in average and "
            f"{epoch_timer.median_epoch_time():.2f}s in median."
        )
        logger.info(
            f"For epoch {cur_epoch}, each iteraction takes "
            f"{epoch_timer.last_epoch_time()/len(train_loaders[0]):.2f}s in average. "
            f"From epoch {start_epoch} to {cur_epoch}, each iteraction takes "
            f"{epoch_timer.avg_epoch_time()/len(train_loaders[0]):.2f}s in average."
        )

        is_checkp_epoch = cu.is_checkpoint_epoch(
            cfg,
            cur_epoch,
            None
        )
        is_eval_epoch = misc.is_eval_epoch(
            cfg,
            cur_epoch,
            None
        )

        # Compute precise BN stats.
        if (
            (is_checkp_epoch or is_eval_epoch)
            and cfg.BN.USE_PRECISE_STATS
            and len(get_bn_modules(model)) > 0
        ):
            calculate_and_update_precise_bn(
                precise_bn_loader,
                model,
                min(cfg.BN.NUM_BATCHES_PRECISE, len(precise_bn_loader)),
                cfg.NUM_GPUS > 0,
            )
        _ = misc.aggregate_sub_bn_stats(model)

        # Save a checkpoint.
        if is_checkp_epoch:
            cu.save_checkpoint(
                cfg.OUTPUT_DIR,
                model,
                optimizer_f,
                cur_epoch,
                cfg,
                scaler if cfg.TRAIN.MIXED_PRECISION else None,
            )
        # Evaluate the model on validation set.
        if is_eval_epoch:
            eval_epoch(
                val_loader,
                model,
                val_meter,
                cur_epoch,
                cfg,
                writer,
            )

    if writer is not None:
        writer.close()
    raise SystemExit('Training Ends')


from datetime import datetime

import pandas as pd
import pydash as _
from bs4 import BeautifulSoup

from Base import NSEBase

class NSE(NSEBase):
    """
        A class to interact with NSE (National Stock Exchange) API.

        Attributes:
            valid_pcr_fields : list of valid fields for put-call ratio calculation

        Methods:
            __init__ : Initialize the NSE class
            get_option_chain : Get the option chain for a given ticker
            get_raw_option_chain : Get the raw option chain data for a given ticker
            get_options_expiry : Get the next expiry date for a given ticker
            get_all_derivatives_enabled_stocks : Get the list of equities available for derivatives trading
            get_equity_future_trade_info : Get the trade information of active future contracts for a given ticker
            get_equity_options_trade_info : Get the trade information of equity options for a given ticker
            _mapped_index_ticker_for_futures : Get the mapped index ticker for index futures
            get_index_futures_data : Get the data for index futures of a given index or ticker
            get_currency_futures : Get the data for currency futures
            get_commodity_futures : Get the data for commodity futures
            get_pcr : Get the put-call ratio for a given ticker and expiry date
    """

    def __init__(self) -> None:
        """
            The __init__ function is called when the class is instantiated.
            It sets up the session and headers for all subsequent requests.

            :param self: Represent the instance of the class

            :return: Nothing
        """

        super().__init__()
        self.valid_pcr_fields = ['oi', 'volume']

    # ----------------------------------------------------------------------------------------------------------------
    # Utility Functions

    def get_option_chain(self, ticker: str, is_index: bool = True, expiry: datetime = None) -> pd.DataFrame:
        """
            The get_option_chain function takes a ticker as input and returns the option chain for that ticker. The
            function uses the try_n_times_get_response function to get a response from NSE's API, which is then converted
            into a DataFrame using pd.json_normalize.

            :param self: Represent the instance of the class
            :param ticker: Specify the stock ticker for which we want to get the option chain its also called symbol in
            NSE
            :param is_index: (optional) Boolean value Specifies the given ticker is an index or not
            :param expiry: (optional) It takes the `expiry date` in the datetime format of the options contracts,
            default is very next expiry day

            :return: A dataframe with option chain
        """

        params = {'symbol': ticker}

        if is_index:
            url = f'{self._base_url}/api/option-chain-indices'
        else:
            url = f'{self._base_url}/api/option-chain-equities'

        response = self.hit_and_get_data(url, params=params)
        if expiry is None:
            df = pd.DataFrame(pd.json_normalize(_.get(response, 'filtered.data', {}), sep='_')).set_index('strikePrice')
        else:
            df = pd.DataFrame(pd.json_normalize(_.get(response, 'records.data', {}), sep='_')).set_index('strikePrice')
            df = df[df['expiryDate'] == expiry.strftime('%d-%b-%Y')]
        return df

    def get_raw_option_chain(self, ticker: str, is_index: bool = True) -> dict:
        """
            The get_option_chain function takes a ticker as input and returns the option chain for that ticker.
            The function uses the try_n_times_get_response function to get a response from NSE's API, which is
            then converted into a DataFrame using pd.json_normalize.

            :param is_index: Boolean value Specifies the given ticker is an index or not
            :param self: Represent the instance of the class
            :param ticker: Specify the stock ticker for which we want to get the option chain

            :return: A dataframe with option chain data
        """

        params = {'symbol': ticker}

        if is_index:
            url = f'{self._base_url}/api/option-chain-indices'
        else:
            url = f'{self._base_url}/api/option-chain-equities'

        response = self.hit_and_get_data(url, params=params)
        return response

    def get_options_expiry(self, ticker: str, is_index: bool = False) -> datetime:
        """
            The get_expiry function takes in a ticker and returns the next expiry date for that ticker.
            The function uses the NSE API to get all expiry dates for a given ticker, sorts them in ascending order,
            and then returns the nth element of this sorted list.

            :param self: Represent the instance of the class
            :param ticker: Specify the ticker / symbol for which we want to get the expiry date
            :param is_index: Boolean value Specifies the given ticker is an index or not

            :return: The very next expiry date
        """

        params = {'symbol': ticker}
        if is_index:
            url = f'{self._base_url}/api/option-chain-indices'
        else:
            url = f'{self._base_url}/api/option-chain-equities'
        response = self.hit_and_get_data(url, params=params)
        dates = sorted([datetime.strptime(date_str, "%d-%b-%Y") for date_str in
                        response.get('records', {}).get('expiryDates', [])])
        return dates

    # ----------------------------------------------------------------------------------------------------------------_
    # Equity Futures
    def get_all_derivatives_enabled_stocks(self) -> list:
        """
            The get_all_derivatives_enabled_stocks provides the list of Equities available for derivative trading

            :param self: Represent the instance of the class

            :return: List of all Equities tickers / symbols for which derivative trading is allowed
        """

        response = self.hit_and_get_data(f'{self._base_url}/api/master-quote')
        return response

    def get_equity_future_trade_info(self, ticker: str) -> pd.DataFrame:
        """
            The get_equity_future_trade_info provides all active future contracts trade information including its price
            details

            :param self: Represent the instance of the class
            :param ticker: Specify the ticker / symbol for which we want to get the expiry date

            :return: A DataFrame of trade info data of Equity Future contracts
        """

        params = {'symbol': ticker}
        response = self.hit_and_get_data(f'{self._base_url}/api/quote-derivative', params=params)
        future_data = []
        for fno_data in response.get('stocks', []):
            if fno_data.get('metadata', {}).get('instrumentType') == 'Stock Futures':
                future_data.append(fno_data)

        df = pd.DataFrame(pd.json_normalize(future_data, sep='_'))
        df['ticker'] = response.get('info', {}).get('symbol', '')
        df['companyName'] = response.get('info', {}).get('companyName', '')
        df['industry'] = response.get('info', {}).get('industry', '')
        df['fut_timestamp'] = response.get('fut_timestamp', '')
        return df

    # ----------------------------------------------------------------------------------------------------------------
    # Equity Options

    def get_equity_options_trade_info(self, ticker: str) -> pd.DataFrame:
        """
            Gets equity options trade information for a given ticker.

            :param ticker: Ticker symbol of the equity options trade.

            :return: DataFrame containing the trade information.
        """

        params = {'symbol': ticker}
        response = self.hit_and_get_data(f'{self._base_url}/api/quote-derivative', params=params)
        future_data = []
        for fno_data in response.get('stocks', []):
            if fno_data.get('metadata', {}).get('instrumentType') == 'Stock Options':
                future_data.append(fno_data)

        df = pd.DataFrame(pd.json_normalize(future_data, sep='_'))
        df['ticker'] = response.get('info', {}).get('symbol', '')
        df['companyName'] = response.get('info', {}).get('companyName', '')
        df['industry'] = response.get('info', {}).get('industry', '')
        df['opt_timestamp'] = response.get('opt_timestamp', '')
        return df

    # ----------------------------------------------------------------------------------------------------------------
    # Index Futures
    def _mapped_index_ticker_for_futures(self) -> dict:
        """
            Mapped index ticker will give dict of available options with its corresponding ticker value

            :param self: Represent the instance of the class

            :return: A dict obj with all FUTURES mappings
        """

        response = self.session.get(f'{self._base_url}//market-data/equity-derivatives-watch',
                                    headers=self.headers)
        soup = BeautifulSoup(response.text, features="html5lib")
        all_derivative_options = soup.find_all('option', attrs={"rel": "derivative"})
        mapped_index_ticker = {}
        for i in all_derivative_options:
            mapped_index_ticker[i.get_text().lower()] = i['value']

        return mapped_index_ticker

    def get_index_futures_data(self, index_or_ticker: str) -> pd.DataFrame:
        """
            Fetches index futures data.

            :param self: Represent the instance of the class
            :param index_or_ticker: Name or ticker symbol of the index.

            :return: DataFrame containing the FUTURES data
        """

        index_or_ticker = index_or_ticker.lower()
        mapped_tickers = {}

        try:
            mapped_tickers = self._mapped_index_ticker_for_futures()
        except Exception as err:
            print(
                f'Exception in fetching mapped ticker for this index try to pass actual ticker in the next call,  '
                f'Exact error : {err}')

        if index_or_ticker in mapped_tickers.keys():
            ticker_to_used = mapped_tickers[index_or_ticker]
        else:
            ticker_to_used = index_or_ticker
        params = {'index': ticker_to_used}
        response = self.hit_and_get_data(f'{self._base_url}/api/liveEquity-derivatives', params=params)
        df = pd.DataFrame(response.get('data', []))
        return df

    # ----------------------------------------------------------------------------------------------------------------
    # Currency

    def get_currency_futures(self) -> pd.DataFrame:
        """
            Fetches currency futures data.

            :param self: Represent the instance of the class

            :return: DataFrame containing the currency futures data
        """

        params = {'index': 'live_market_currency', 'key': 'INR'}
        response = self.hit_and_get_data(
            f'{self._base_url}/api/liveCurrency-derivatives', params=params)
        df = pd.DataFrame(response.get('data', []))
        return df

    # ----------------------------------------------------------------------------------------------------------------
    # Commodity
    def get_commodity_futures(self) -> pd.DataFrame:
        """
            Fetches commodity futures data.

            :param self: Represent the instance of the class

            :return: Pd.DataFrame: DataFrame containing the currency futures data
        """
        response = self.hit_and_get_data(f'{self._base_url}/api/liveCommodity-derivatives')
        df = pd.DataFrame(response.get('data', []))
        return df

    def get_pcr(self, ticker: str, is_index: bool = True, on_field: str = 'OI', expiry: datetime = None) -> float:
        """
            Calculate the put-call ratio (PCR) for a given ticker.

            :param self: Represent the instance of the class
            :param ticker: The ticker symbol.
            :param is_index: Boolean value Specifies the given ticker is an index or not
            :param expiry: The expiry date of the option contract. Defaults to None.
            :param on_field: The field to calculate PCR on. `Volume` or `oi` (open-interest) Default to 'OI'.

            :return: The calculated PCR value
        """

        on_field = on_field.lower()
        if on_field not in self.valid_pcr_fields:
            print(f'Un-supported filed is passed only these are the fields available : {self.valid_pcr_fields}')
            return 0

        if expiry is None:
            df = self.get_option_chain(ticker, is_index=is_index)
        else:
            df = self.get_option_chain(ticker, is_index=is_index, expiry=expiry)

        if df.shape[0] == 0:
            print('Your filters lead to empty DataSet check all params, expiry, etc; returning 0 as default')
            return 0
        if on_field == 'oi':
            put_oi = df['PE_openInterest'].sum()
            call_oi = df['CE_openInterest'].sum()
            return put_oi / call_oi
        else:
            put_vol = df['PE_totalTradedVolume'].sum()
            call_vol = df['CE_totalTradedVolume'].sum()
            return put_vol / call_vol