File size: 78,636 Bytes
2b91026
 
 
 
 
 
0abefc5
 
2b91026
 
 
 
0abefc5
2b91026
 
 
 
 
 
 
a14618a
 
 
2b91026
9e66ab7
57bed38
2b91026
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
9bac4cf
2b91026
 
 
9bac4cf
2b91026
 
 
 
 
 
9bac4cf
2b91026
 
 
9bac4cf
2b91026
 
 
 
 
a14618a
 
 
0301e99
84a3e25
2b91026
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
a14618a
2b91026
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
a14618a
 
2b91026
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633
634
635
636
637
638
639
640
641
642
643
644
645
646
647
648
649
650
651
652
653
654
655
656
657
658
659
660
661
662
663
664
665
666
667
668
669
670
671
672
673
674
675
676
677
678
679
680
681
682
683
684
685
686
687
688
689
690
691
692
693
694
695
696
697
698
699
700
701
702
703
704
705
706
707
708
709
710
711
712
713
714
715
716
717
718
719
720
721
722
723
724
725
726
727
728
729
730
731
732
733
734
735
736
737
738
739
740
741
742
743
744
745
746
747
748
749
750
751
752
753
754
755
756
757
758
759
760
761
762
763
764
765
766
767
768
769
770
771
772
773
774
775
776
777
778
779
780
781
782
783
784
785
786
787
788
789
790
791
792
793
794
795
796
797
798
799
800
801
802
803
804
805
806
807
808
809
810
811
812
813
814
815
816
817
818
819
820
821
822
823
824
825
826
827
828
829
830
831
832
833
834
835
836
837
838
839
840
841
842
843
844
845
846
847
848
849
850
851
852
853
854
855
856
857
858
859
860
861
862
863
864
865
866
867
868
869
870
871
872
873
874
875
876
877
878
879
880
881
882
883
884
885
886
887
888
889
890
891
892
893
894
895
896
897
898
899
900
901
902
903
904
905
906
907
908
909
910
911
912
913
914
915
916
917
918
919
920
921
922
923
924
925
926
927
928
929
930
931
932
933
934
935
936
937
938
939
940
941
942
943
944
945
946
947
948
949
950
951
952
953
954
955
956
957
958
959
960
961
962
963
964
965
966
967
968
969
970
971
972
973
974
975
976
977
978
979
980
981
982
983
984
985
986
987
988
989
990
991
992
993
994
995
996
997
998
999
1000
1001
1002
1003
1004
1005
1006
1007
1008
1009
1010
1011
1012
1013
1014
1015
1016
1017
1018
1019
1020
1021
1022
1023
1024
1025
1026
1027
1028
1029
1030
1031
1032
1033
1034
1035
1036
1037
1038
1039
1040
1041
1042
1043
1044
1045
1046
1047
1048
1049
1050
1051
1052
1053
1054
1055
1056
1057
1058
1059
1060
1061
1062
1063
1064
1065
1066
1067
1068
1069
1070
1071
1072
1073
1074
1075
1076
1077
1078
1079
1080
1081
1082
1083
1084
1085
1086
1087
1088
1089
1090
1091
1092
1093
1094
1095
1096
1097
1098
1099
1100
1101
1102
1103
1104
1105
1106
1107
1108
1109
1110
1111
1112
1113
1114
1115
1116
1117
1118
1119
1120
1121
1122
1123
1124
1125
1126
1127
1128
1129
1130
1131
1132
1133
1134
1135
1136
1137
1138
1139
1140
1141
1142
1143
1144
1145
1146
1147
1148
1149
1150
1151
1152
1153
1154
1155
1156
1157
1158
1159
1160
1161
1162
1163
1164
1165
1166
1167
1168
1169
1170
1171
1172
1173
1174
1175
1176
1177
1178
1179
1180
1181
1182
1183
1184
1185
1186
1187
1188
1189
1190
1191
1192
1193
1194
1195
1196
1197
1198
1199
1200
1201
1202
1203
1204
1205
1206
1207
1208
1209
1210
1211
1212
1213
1214
1215
1216
1217
1218
1219
1220
1221
1222
1223
1224
1225
1226
1227
1228
1229
1230
1231
1232
1233
1234
1235
1236
1237
1238
1239
1240
1241
1242
1243
1244
1245
1246
1247
1248
1249
1250
1251
1252
1253
1254
1255
1256
1257
1258
1259
1260
1261
1262
1263
1264
1265
1266
1267
1268
1269
1270
1271
1272
1273
1274
1275
1276
1277
1278
1279
1280
1281
1282
1283
1284
1285
1286
1287
1288
1289
1290
1291
1292
1293
1294
1295
1296
1297
1298
1299
1300
1301
1302
1303
1304
1305
1306
1307
1308
1309
1310
1311
1312
1313
1314
1315
1316
1317
1318
1319
1320
1321
1322
1323
1324
1325
1326
1327
1328
1329
1330
1331
1332
1333
1334
1335
1336
1337
1338
1339
1340
1341
1342
1343
1344
1345
1346
1347
1348
1349
1350
1351
1352
1353
1354
1355
1356
1357
1358
1359
1360
1361
1362
1363
1364
1365
1366
1367
1368
1369
1370
1371
1372
1373
1374
1375
1376
1377
1378
1379
1380
1381
1382
1383
1384
1385
1386
1387
1388
1389
1390
1391
1392
1393
1394
1395
1396
1397
1398
1399
1400
1401
1402
1403
1404
1405
1406
1407
1408
1409
1410
1411
1412
1413
1414
1415
1416
1417
1418
1419
1420
1421
1422
1423
1424
1425
1426
1427
1428
1429
1430
1431
1432
1433
1434
1435
1436
1437
1438
1439
1440
1441
1442
1443
1444
1445
1446
1447
1448
1449
1450
1451
1452
1453
1454
1455
1456
1457
1458
1459
1460
1461
1462
1463
1464
1465
1466
1467
1468
1469
1470
1471
1472
1473
1474
1475
1476
1477
1478
1479
1480
1481
1482
1483
1484
1485
1486
1487
1488
1489
1490
1491
1492
1493
1494
1495
1496
1497
1498
1499
1500
1501
1502
1503
1504
1505
1506
1507
1508
1509
1510
1511
1512
1513
1514
1515
1516
1517
1518
1519
1520
1521
1522
1523
1524
1525
1526
1527
1528
1529
1530
1531
1532
1533
1534
1535
1536
1537
1538
1539
1540
1541
1542
1543
1544
1545
1546
1547
1548
1549
1550
1551
1552
1553
1554
1555
1556
1557
1558
1559
1560
1561
1562
1563
1564
1565
1566
1567
1568
1569
1570
1571
1572
1573
1574
1575
1576
1577
1578
1579
1580
1581
1582
1583
1584
1585
1586
1587
1588
1589
1590
1591
1592
1593
1594
1595
1596
1597
1598
1599
1600
1601
1602
1603
1604
1605
1606
1607
1608
1609
1610
1611
1612
1613
1614
1615
1616
1617
1618
1619
1620
1621
1622
1623
1624
1625
1626
1627
1628
1629
1630
1631
1632
1633
1634
1635
1636
1637
1638
1639
1640
1641
1642
1643
1644
1645
1646
1647
1648
1649
1650
1651
1652
1653
1654
1655
1656
1657
1658
1659
1660
1661
1662
1663
1664
1665
1666
1667
1668
1669
1670
1671
1672
1673
1674
1675
1676
1677
1678
1679
1680
1681
1682
1683
1684
1685
1686
1687
1688
1689
1690
1691
1692
1693
1694
1695
1696
1697
1698
1699
1700
1701
1702
1703
1704
1705
1706
1707
1708
1709
1710
1711
1712
1713
1714
1715
1716
1717
1718
1719
1720
1721
1722
1723
1724
1725
1726
1727
1728
1729
1730
1731
1732
1733
1734
1735
1736
1737
1738
1739
1740
1741
1742
1743
1744
1745
1746
1747
1748
1749
1750
1751
1752
1753
1754
1755
1756
1757
1758
1759
1760
1761
1762
1763
1764
1765
1766
1767
1768
1769
1770
1771
1772
1773
1774
1775
1776
1777
1778
1779
1780
1781
1782
1783
1784
1785
1786
1787
1788
1789
1790
1791
1792
1793
1794
1795
1796
1797
1798
1799
1800
1801
1802
1803
1804
1805
1806
1807
1808
1809
1810
1811
1812
1813
1814
1815
1816
1817
1818
1819
1820
1821
1822
1823
1824
1825
1826
1827
1828
1829
1830
1831
1832
1833
1834
1835
1836
1837
1838
1839
1840
1841
1842
1843
1844
1845
1846
1847
1848
1849
1850
1851
1852
1853
1854
1855
1856
1857
1858
1859
1860
1861
1862
1863
1864
1865
1866
1867
1868
1869
1870
1871
1872
1873
1874
1875
1876
1877
1878
1879
1880
1881
1882
1883
1884
1885
1886
1887
1888
1889
1890
1891
1892
1893
1894
1895
1896
1897
1898
1899
1900
1901
1902
1903
1904
1905
1906
1907
1908
1909
1910
1911
1912
1913
1914
1915
1916
1917
1918
1919
1920
1921
1922
1923
1924
1925
1926
1927
1928
1929
1930
1931
import tushare as ts
import matplotlib.pyplot as plt
import pandas as pd
import os
import random
from matplotlib.ticker import MaxNLocator
#from prettytable import PrettyTable
#from blessed import Terminal
import time
from datetime import datetime, timedelta
import numpy as np
import mplfinance as mpf

from typing import Optional
import matplotlib.font_manager as fm
from matplotlib.lines import Line2D
from typing import Union, Any
from sklearn.linear_model import LinearRegression


# plt.rcParams['font.sans-serif'] = ['Arial Unicode MS']
# plt.rcParams['axes.unicode_minus'] = False


font_path = './fonts/SimHei.ttf'
font_prop = fm.FontProperties()


tushare_token = os.getenv('TUSHARE_TOKEN')
pro = ts.pro_api(tushare_token)

# def last_month_end(date_str:str=''):
#     date_obj = datetime.strptime(date_str, '%Y%m%d')
#     current_month = date_obj.month
#     current_year = date_obj.year
#
#     if current_month == 1:
#         last_month = 12
#         last_year = current_year - 1
#     else:
#         last_month = current_month - 1
#         last_year = current_year
#
#     if date_obj.month != (date_obj + timedelta(days=1)).month:
#         last_month_end_date = date_obj
#     else:
#         last_day_of_last_month = (date_obj.replace(day=1) - timedelta(days=1)).day
#         last_month_end_date = datetime(last_year, last_month, last_day_of_last_month)
#
#     return last_month_end_date.strftime('%Y%m%d')



def get_last_year_date(date_str: str = '') -> str:
    """
        This function takes a date string in the format YYYYMMDD and returns the date string one year prior to the input date.

        Args:
        - date_str: string, the input date in the format YYYYMMDD

        Returns:
        - string, the date one year prior to the input date in the format YYYYMMDD
        """
    dt = datetime.strptime(date_str, '%Y%m%d')
    # To calculate the date one year ago
    one_year_ago = dt - timedelta(days=365)

    # To format the date as a string
    one_year_ago_str = one_year_ago.strftime('%Y%m%d')

    return one_year_ago_str


def get_adj_factor(stock_code: str = '', start_date: str = '', end_date: str = '') -> pd.DataFrame:
    # Get stock price adjustment factors. Retrieve the stock price adjustment factors for a single stock's entire historical data or for all stocks on a single trading day.
    # The input includes the stock code, start date, end date, and trading date, all in string format with the date in the YYYYMMDD format
    # The return value is a dataframe containing the stock code, trading date, and adjustment factor
    # ts_code	str	股票代码
    # adj_factor	float	复权因子
    """
       This function retrieves the adjusted stock prices for a given stock code and date range.

       Args:
       - stock_code: string, the stock code to retrieve data for
       - start_date: string, the start date in the format YYYYMMDD
       - end_date: string, the end date in the format YYYYMMDD

       Returns:
       - dataframe, a dataframe containing the stock code, trade date, and adjusted factor

       This will retrieve the adjusted stock prices for the stock with code '000001.SZ' between the dates '20220101' and '20220501'.
       """
    df = pro.adj_factor(**{
        "ts_code": stock_code,
        "trade_date": "",
        "start_date": start_date,
        "end_date": end_date,
        "limit": "",
        "offset": ""
    }, fields=[
        "ts_code",
        "trade_date",
        "adj_factor"
    ])

    return df

def get_stock_code(stock_name: str) -> str:
    # Retrieve the stock code of a given stock name. If we call get_stock_code('贵州茅台'), it will return '600519.SH'.


    df = pd.read_csv('tushare_stock_basic_20230421210721.csv')
    try:
        code = df.loc[df.name==stock_name].ts_code.iloc[0]
        return code
    except:
        return None




def get_stock_name_from_code(stock_code: str) -> str:
    """
        Reads a local file to retrieve the stock name from a given stock code.

        Args:
        - stock_code (str): The code of the stock.

        Returns:
        - str: The stock name of the given stock code.
        """
    # For example,if we call get_stock_name_from_code('600519.SH'), it will return '贵州茅台'.


    df = pd.read_csv('tushare_stock_basic_20230421210721.csv')
    name = df.loc[df.ts_code == stock_code].name.iloc[0]

    return name

def get_stock_prices_data(stock_name: str='', start_date: str='', end_date: str='', freq:str='daily') -> pd.DataFrame:
    """
        Retrieves the daily/weekly/monthly price data for a given stock code during a specific time period. get_stock_prices_data('贵州茅台','20200120','20220222','daily')

        Args:
        - stock_name (str)
        - start_date (str): The start date in the format 'YYYYMMDD'.
        - end_date (str): The end date in 'YYYYMMDD'.
        - freq (str): The frequency of the price data, can be 'daily', 'weekly', or 'monthly'.

        Returns:
        - pd.DataFrame: A dataframe that contains the daily/weekly/monthly data. The output columns contain stock_code, trade_date, open, high, low, close, pre_close(昨天收盘价), change(涨跌额), pct_chg(涨跌幅),vol(成交量),amount(成交额)
        """

    stock_code = get_stock_code(stock_name)

    if freq == 'daily':
        stock_data = pro.daily(**{
            "ts_code": stock_code,
            "trade_date": '',
            "start_date": start_date,
            "end_date": end_date,
            "offset": "",
            "limit": ""
        }, fields=[
            "ts_code",
            "trade_date",
            "open",
            "high",
            "low",
            "close",
            "pre_close",
            "change",
            "pct_chg",
            "vol",
            "amount"
        ])

    elif freq == 'weekly':
        stock_data = pro.weekly(**{
            "ts_code": stock_code,
            "trade_date": '',
            "start_date": start_date,
            "end_date": end_date,
            "limit": "",
            "offset": ""
        }, fields=[
            "ts_code",
            "trade_date",
            "close",
            "open",
            "high",
            "low",
            "pre_close",
            "change",
            "pct_chg",
            "vol",
            "amount"
        ])
    elif freq == 'monthly':
        stock_data = pro.monthly(**{
            "ts_code": stock_code,
            "trade_date": '',
            "start_date": start_date,
            "end_date": end_date,
            "limit": "",
            "offset": ""
        }, fields=[
            "ts_code",
            "trade_date",
            "close",
            "open",
            "high",
            "low",
            "pre_close",
            "change",
            "pct_chg",
            "vol",
            "amount"
        ])


    adj_f = get_adj_factor(stock_code, start_date, end_date)
    stock_data = pd.merge(stock_data, adj_f, on=['ts_code', 'trade_date'])
    # Multiply the values of open, high, low, and close by their corresponding adjustment factors.
    # To obtain the adjusted close price
    stock_data[['open', 'high', 'low', 'close']] *= stock_data['adj_factor'].values.reshape(-1, 1)

    #stock_data.rename(columns={'vol': 'volume'}, inplace=True)
    df = pd.read_csv('tushare_stock_basic_20230421210721.csv')
    stock_data_merged = pd.merge(stock_data, df, on='ts_code')
    stock_data_merged.rename(columns={'ts_code': 'stock_code'}, inplace=True)
    stock_data_merged.rename(columns={'name': 'stock_name'}, inplace=True)
    stock_data_merged = stock_data_merged.sort_values(by='trade_date', ascending=True)  # To sort the DataFrame by date in ascending order
    return stock_data_merged



def get_stock_technical_data(stock_name: str, start_date: str, end_date: str) -> pd.DataFrame:
    """
        Retrieves the daily technical data of a stock including macd turnover rate, volume, PE ratio, etc. Those technical indicators are usually plotted as subplots in a k-line chart.

        Args:
            stock_name (str):
            start_date (str): Start date "YYYYMMDD"
            end_date (str): End date "YYYYMMDD"

        Returns:
            pd.DataFrame: A DataFrame containing the technical data of the stock,
            including various indicators such as ts_code, trade_date, close, macd_dif, macd_dea, macd, kdj_k, kdj_d, kdj_j, rsi_6, rsi_12, boll_upper, boll_mid, boll_lower, cci, turnover_rate, turnover_rate_f, volume_ratio, pe_ttm(市盈率), pb(市净率), ps_ttm, dv_ttm, total_share, float_share, free_share, total_mv, circ_mv

    """

    # Technical factors
    stock_code = get_stock_code(stock_name)
    stock_data1 = pro.stk_factor(**{
        "ts_code": stock_code,
        "start_date": start_date,
        "end_date": end_date,
        "trade_date": '',
        "limit": "",
        "offset": ""
    }, fields=[
        "ts_code",
        "trade_date",
        "close",
        "macd_dif",
        "macd_dea",
        "macd",
        "kdj_k",
        "kdj_d",
        "kdj_j",
        "rsi_6",
        "rsi_12",
        "rsi_24",
        "boll_upper",
        "boll_mid",
        "boll_lower",
        "cci"
    ])
    # Trading factors
    stock_data2 = pro.daily_basic(**{
        "ts_code": stock_code,
        "trade_date": '',
        "start_date": start_date,
        "end_date": end_date,
        "limit": "",
        "offset": ""
    }, fields=[
        "ts_code",  #
        "trade_date",
        "turnover_rate",
        "turnover_rate_f",
        "volume_ratio",
        "pe_ttm",
        "pb",
        "ps_ttm",
        "dv_ttm",
        "total_share",
        "float_share",
        "free_share",
        "total_mv",
        "circ_mv"
    ])

    #
    stock_data = pd.merge(stock_data1, stock_data2, on=['ts_code', 'trade_date'])
    df = pd.read_csv('tushare_stock_basic_20230421210721.csv')
    stock_data_merged = pd.merge(stock_data, df, on='ts_code')
    stock_data_merged = stock_data_merged.sort_values(by='trade_date', ascending=True)

    stock_data_merged.drop(['symbol'], axis=1, inplace=True)

    stock_data_merged.rename(columns={'ts_code': 'stock_code'}, inplace=True)
    stock_data_merged.rename(columns={'name': 'stock_name'}, inplace=True)

    return stock_data_merged


def plot_stock_data(stock_data: pd.DataFrame, ax: Optional[plt.Axes] = None, figure_type: str = 'line', title_name: str ='') -> plt.Axes:

    """
    This function plots stock data.

    Args:
    - stock_data: pandas DataFrame, the stock data to plot. The DataFrame should contain three columns:
        - Column 1: trade date in 'YYYYMMDD'
        - Column 2: Stock name or code (string format)
        - Column 3: Index value (numeric format)
        The DataFrame can be time series data or cross-sectional data. If it is time-series data, the first column represents different trade time, the second column represents the same name. For cross-sectional data, the first column is the same, the second column contains different stocks.

    - ax: matplotlib Axes object, the axes to plot the data on
    - figure_type: the type of figure (either 'line' or 'bar')
    - title_name

    Returns:
    - matplotlib Axes object, the axes containing the plot
    """

    index_name = stock_data.columns[2]
    name_list = stock_data.iloc[:,1]
    date_list = stock_data.iloc[:,0]
    if name_list.nunique() == 1 and date_list.nunique() != 1:
        # Time Series Data
        unchanged_var = name_list.iloc[0]   # stock name
        x_dim = date_list                   # tradingdate
        x_name = stock_data.columns[0]

    elif name_list.nunique() != 1 and date_list.nunique() == 1:
        # Cross-sectional Data
        unchanged_var = date_list.iloc[0]    # tradingdate
        x_dim = name_list                    # stock name
        x_name = stock_data.columns[1]

        data_size = x_dim.shape[0]



    start_x_dim, end_x_dim = x_dim.iloc[0], x_dim.iloc[-1]

    start_y = stock_data.iloc[0, 2]
    end_y = stock_data.iloc[-1, 2]


    def generate_random_color():
        r = random.randint(0, 255)/ 255.0
        g = random.randint(0, 100)/ 255.0
        b = random.randint(0, 255)/ 255.0
        return (r, g, b)

    color = generate_random_color()
    if ax is None:
        _, ax = plt.subplots()

    if figure_type =='line':
        #

        ax.plot(x_dim, stock_data.iloc[:, 2], label = unchanged_var+'_' + index_name, color=color,linewidth=3)
        #
        plt.scatter(x_dim, stock_data.iloc[:, 2], color=color,s=3)  # Add markers to the data points

        #
        #ax.scatter(x_dim, stock_data.iloc[:, 2],label = unchanged_var+'_' + index_name, color=color, s=3)
        #

        ax.annotate(unchanged_var + ':' + str(round(start_y, 2)) + ' @' + start_x_dim, xy=(start_x_dim, start_y),
                    xytext=(start_x_dim, start_y),
                    textcoords='data', fontsize=14,color=color, horizontalalignment='right',fontproperties=font_prop)

        ax.annotate(unchanged_var + ':' + str(round(end_y, 2)) +' @' + end_x_dim, xy=(end_x_dim, end_y),
                    xytext=(end_x_dim, end_y),
                    textcoords='data', fontsize=14, color=color, horizontalalignment='left',fontproperties=font_prop)


    elif figure_type == 'bar':
        ax.bar(x_dim, stock_data.iloc[:, 2], label = unchanged_var + '_' + index_name, width=0.3, color=color)
        ax.annotate(unchanged_var + ':' + str(round(start_y, 2)) + ' @' + start_x_dim, xy=(start_x_dim, start_y),
                    xytext=(start_x_dim, start_y),
                    textcoords='data', fontsize=14, color=color, horizontalalignment='right',fontproperties=font_prop)

        ax.annotate(unchanged_var + ':' + str(round(end_y, 2)) + ' @' + end_x_dim, xy=(end_x_dim, end_y),
                    xytext=(end_x_dim, end_y),
                    textcoords='data', fontsize=14, color=color, horizontalalignment='left',fontproperties=font_prop)

    plt.xticks(x_dim,rotation=45)                                                  #
    ax.xaxis.set_major_locator(MaxNLocator( integer=True, prune=None, nbins=100))  #


    plt.xlabel(x_name, fontproperties=font_prop,fontsize=18)
    plt.ylabel(f'{index_name}', fontproperties=font_prop,fontsize=16)
    ax.set_title(title_name , fontproperties=font_prop,fontsize=16)
    plt.legend(prop=font_prop)  # 显示图例
    fig = plt.gcf()
    fig.set_size_inches(18, 12)

    return ax


def query_fund_Manager(Manager_name: str) -> pd.DataFrame:
    # 代码fund_code,公告日期ann_date,基金经理名字name,性别gender,出生年份birth_year,学历edu,国籍nationality,开始管理日期begin_date,结束日期end_date,简历resume
    """
        Retrieves information about a fund manager.

        Args:
            Manager_name (str): The name of the fund manager.

        Returns:
            df (DataFrame): A DataFrame containing the fund manager's information, including the fund codes, announcement dates,
                            manager's name, gender, birth year, education, nationality, start and end dates of managing funds,
                            and the manager's resume.
    """

    df = pro.fund_manager(**{
        "ts_code": "",
        "ann_date": "",
        "name": Manager_name,
        "offset": "",
        "limit": ""
    }, fields=[
        "ts_code",
        "ann_date",
        "name",
        "gender",
        "birth_year",
        "edu",
        "nationality",
        "begin_date",
        "end_date",
        "resume"
    ])
    #
    df.rename(columns={'ts_code': 'fund_code'}, inplace=True)
    # To query the fund name based on the fund code and store it in a new column called fund_name, while removing the rows where the fund name is not found
    df['fund_name'] = df['fund_code'].apply(lambda x: query_fund_name_or_code('', x))
    df.dropna(subset=['fund_name'], inplace=True)
    df.rename(columns={'name': 'manager_name'}, inplace=True)
    #
    df_out = df[['fund_name','fund_code','ann_date','manager_name','begin_date','end_date']]

    return df_out


# def save_stock_prices_to_csv(stock_prices: pd.DataFrame, stock_name: str, file_path: str) -> None:
#
#     """
#         Saves the price data of a specific stock symbol during a specific time period to a local CSV file.
#
#         Args:
#         - stock_prices (pd.DataFrame): A pandas dataframe that contains the daily price data for the given stock symbol during the specified time period.
#         - stock_name (str): The name of the stock.
#         - file_path (str): The file path where the CSV file will be saved.
#
#         Returns:
#         - None: The function only saves the CSV file to the specified file path.
#     """
#     # The function checks if the directory to save the CSV file exists and creates it if it does not exist.
#     # The function then saves the price data of the specified stock symbol during the specified time period to a local CSV file with the name {stock_name}_price_data.csv in the specified file path.
#
#
#     if not os.path.exists(file_path):
#         os.makedirs(file_path)
#
#
#     file_path = f"{file_path}{stock_name}_stock_prices.csv"
#     stock_prices.to_csv(file_path, index_label='Date')
#     print(f"Stock prices for {stock_name} saved to {file_path}")


def calculate_stock_index(stock_data: pd.DataFrame, index:str='close') -> pd.DataFrame:
    """
        Calculate a specific index of a stock based on its price information.

        Args:
            stock_data (pd.DataFrame): DataFrame containing the stock's price information.
            index (str, optional): The index to calculate. The available options depend on the column names in the
                input stock price data. Additionally, there are two special indices: 'candle_K' and 'Cumulative_Earnings_Rate'.

        Returns:
            DataFrame containing the corresponding index data of the stock. In general, it includes three columns: 'trade_date', 'name', and the corresponding index value.
            Besides, if index is 'candle_K', the function returns the DataFrame containing 'trade_date', 'Open', 'High', 'Low', 'Close', 'Volume','name' column.
            If index is a technical index such as 'macd' or a trading index likes 'pe_ttm', the function returns the DataFrame with corresponding columns.
        """


    if 'stock_name' not in  stock_data.columns and 'index_name' in stock_data.columns:
        stock_data.rename(columns={'index_name': 'stock_name'}, inplace=True)
    #
    index = index.lower()
    if index=='Cumulative_Earnings_Rate' or index =='Cumulative_Earnings_Rate'.lower() :
        stock_data[index] = (1 + stock_data['pct_chg'] / 100.).cumprod() - 1.
        stock_data[index] = stock_data[index] * 100.
        if 'stock_name' in stock_data.columns :
           selected_index = stock_data[['trade_date', 'stock_name', index]].copy()
        #
        if 'fund_name' in stock_data.columns:
            selected_index = stock_data[['trade_date', 'fund_name', index]].copy()
        return selected_index

    elif index == 'candle_K' or index == 'candle_K'.lower():
        #tech_df = tech_df.drop(['name', 'symbol', 'industry', 'area','market','list_date','ts_code','close'], axis=1)
        # Merge two DataFrames based on the 'trade_date' column.

        stock_data = stock_data.rename(
            columns={'open': 'Open', 'high': 'High', 'low': 'Low', 'close': 'Close',
                     'vol': 'Volume'})
        selected_index = stock_data[['trade_date', 'Open', 'High', 'Low', 'Close', 'Volume','stock_name']].copy()
        return selected_index

    elif index =='macd':
        selected_index = stock_data[['trade_date','macd','macd_dea','macd_dif']].copy()
        return selected_index

    elif index =='rsi':
        selected_index = stock_data[['trade_date','rsi_6','rsi_12']].copy()
        return selected_index

    elif index =='boll':
        selected_index = stock_data[['trade_date', 'boll_upper', 'boll_lower','boll_mid']].copy()
        return selected_index

    elif index =='kdj':
        selected_index = stock_data[['trade_date', 'kdj_k', 'kdj_d','kdj_j']].copy()
        return selected_index

    elif index =='cci':
        selected_index = stock_data[['trade_date', 'cci']].copy()
        return selected_index

    elif index == '换手率':
        selected_index = stock_data[['trade_date', 'turnover_rate','turnover_rate_f']].copy()
        return selected_index

    elif index == '市值':
        selected_index = stock_data[['trade_date', 'total_mv','circ_mv']].copy()
        return selected_index


    elif index in stock_data.columns:
        stock_data = stock_data

        if 'stock_name' in stock_data.columns :
           selected_index = stock_data[['trade_date', 'stock_name', index]].copy()

        if 'fund_name' in stock_data.columns:
            selected_index = stock_data[['trade_date', 'fund_name', index]].copy()
        # Except for candlestick chart and technical indicators, the remaining outputs consist of three columns: date, name, and indicator.
        return selected_index



def rank_index_cross_section(stock_data: pd.DataFrame, Top_k: int = -1, ascending: bool = False) -> pd.DataFrame:
    """
        Sort the cross-sectional data based on the given index.

        Args:
            stock_data : DataFrame containing the cross-sectional data. It should have three columns, and the last column represents the variable to be sorted.
            Top_k : The number of data points to retain after sorting. (Default: -1, which retains all data points)
            ascending: Whether to sort the data in ascending order or not. (Default: False)

        Returns:
            stock_data_selected : DataFrame containing the sorted data. It has the same structure as the input DataFrame.
        """

    index = stock_data.columns[-1]
    stock_data = stock_data.sort_values(by=index, ascending=ascending)
    #stock_data_selected = stock_data[['trade_date','stock_name', index]].copy()
    stock_data_selected = stock_data[:Top_k]
    stock_data_selected = stock_data_selected.drop_duplicates(subset=['stock_name'], keep='first')
    return stock_data_selected


def get_company_info(stock_name: str='') -> pd.DataFrame:
    # ts_code: str	股票代码,  exchange:str	交易所代码SSE上交所 SZSE深交所, chairman:str 法人代表, manager:str 总经理, secretary:str	董秘 # reg_capital:float	注册资本, setup_date:str 注册日期, province:str 所在省份 ,city:str 所在城市
    # introduction:str 公司介绍, website:str 公司主页 , email:str	电子邮件, office:str 办公室 # ann_date: str 公告日期, business_scope:str 经营范围, employees:int	员工人数, main_business:str 主要业务及产品
    """
            This function retrieves company information including stock code, exchange, chairman, manager, secretary,
            registered capital, setup date, province, city, website, email, employees, business scope, main business,
            introduction, office, and announcement date.

            Args:
            - stock_name (str): The name of the stock.

            Returns:
            - pd.DataFrame: A DataFrame that contains the company information.
    """

    stock_code = get_stock_code(stock_name)
    df = pro.stock_company(**{
        "ts_code": stock_code,"exchange": "","status": "", "limit": "","offset": ""
    }, fields=[
        "ts_code","exchange","chairman", "manager","secretary", "reg_capital","setup_date", "province","city",
        "website", "email","employees","business_scope","main_business","introduction","office", "ann_date"
    ])


    en_to_cn = {
        'ts_code': '股票代码',
        'exchange': '交易所代码',
        'chairman': '法人代表',
        'manager': '总经理',
        'secretary': '董秘',
        'reg_capital': '注册资本',
        'setup_date': '注册日期',
        'province': '所在省份',
        'city': '所在城市',
        'introduction': '公司介绍',
        'website': '公司主页',
        'email': '电子邮件',
        'office': '办公室',
        'ann_date': '公告日期',
        'business_scope': '经营范围',
        'employees': '员工人数',
        'main_business': '主要业务及产品'
    }

    df.rename(columns=en_to_cn, inplace=True)
    df.insert(0, '股票名称', stock_name)
    # for column in df.columns:
    #     print(f"[{column}]: {df[column].values[0]}")


    return df





# def get_Financial_data(stock_code: str, report_date: str,  financial_index: str = '' ) -> pd.DataFrame:
#     # report_date的格式为"YYYYMMDD",包括"yyyy0331"为一季报,"yyyy0630"为半年报,"yyyy0930"为三季报,"yyyy1231"为年报
#     # index包含: # current_ratio	流动比率 # quick_ratio	速动比率 # netprofit_margin	销售净利率 # grossprofit_margin	销售毛利率 # roe	净资产收益率 # roe_dt	净资产收益率(扣除非经常损益)
#     # roa	总资产报酬率 # debt_to_assets 资产负债率 # roa_yearly	年化总资产净利率  # q_dtprofit	扣除非经常损益后的单季度净利润 # q_eps	每股收益(单季度)
#     # q_netprofit_margin	销售净利率(单季度) # q_gsprofit_margin	销售毛利率(单季度) # basic_eps_yoy 基本每股收益同比增长率(%) # netprofit_yoy	归属母公司股东的净利润同比增长率(%)   # q_netprofit_yoy	归属母公司股东的净利润同比增长率(%)(单季度) # q_netprofit_qoq	归属母公司股东的净利润环比增长率(%)(单季度) # equity_yoy	净资产同比增长率
#     """
#         Retrieves financial data for a specific stock within a given date range.
#
#         Args:
#             stock_code (str): The stock code or symbol of the company for which financial data is requested.
#             report_date (str): The report date in the format "YYYYMMDD" .
#             financial_index (str, optional): The financial indicator to be queried. If not specified, all available financial
#                                               indicators will be included.
#
#         Returns:
#             pd.DataFrame: A DataFrame containing the financial data for the specified stock and date range. The DataFrame
#                           consists of the following columns: "stock_name",
#                           "trade_date" (reporting period), and the requested financial indicator(s).
#
#         """
#     stock_data = pro.fina_indicator(**{
#         "ts_code": stock_code,
#         "ann_date": "",
#         "start_date": '',
#         "end_date": '',
#         "period": report_date,
#         "update_flag": "1",
#         "limit": "",
#         "offset": ""
#     }, fields=["ts_code","end_date", financial_index])
#
#     stock_name = get_stock_name_from_code(stock_code)
#     stock_data['stock_name'] = stock_name
#     stock_data = stock_data.sort_values(by='end_date', ascending=True)  # 按照日期升序排列
#     # 把end_data列改名为trade_date
#     stock_data.rename(columns={'end_date': 'trade_date'}, inplace=True)
#     stock_financial_data = stock_data[['stock_name', 'trade_date', financial_index]]
#     return stock_financial_data


def get_Financial_data_from_time_range(stock_name:str, start_date:str, end_date:str, financial_index:str='') -> pd.DataFrame:
    # start_date='20190101',end_date='20221231',financial_index='roe', The returned data consists of the ROE values for the entire three-year period from 2019 to 2022.
    # To query quarterly or annual financial report data for a specific moment, "yyyy0331"为一季报,"yyyy0630"为半年报,"yyyy0930"为三季报,"yyyy1231"为年报,例如get_Financial_data_from_time_range("600519.SH", "20190331", "20190331", "roe") means to query the return on equity (ROE) data from the first quarter of 2019,
    #  # current_ratio	流动比率 # quick_ratio	速动比率 # netprofit_margin	销售净利率 # grossprofit_margin	销售毛利率 # roe	净资产收益率 # roe_dt	净资产收益率(扣除非经常损益)
    # roa	总资产报酬率 # debt_to_assets 资产负债率 # roa_yearly	年化总资产净利率  # q_dtprofit	扣除非经常损益后的单季度净利润 # q_eps	每股收益(单季度)
    # q_netprofit_margin	销售净利率(单季度) # q_gsprofit_margin	销售毛利率(单季度) # basic_eps_yoy 基本每股收益同比增长率(%) # netprofit_yoy	归属母公司股东的净利润同比增长率(%)   # q_netprofit_yoy	归属母公司股东的净利润同比增长率(%)(单季度) # q_netprofit_qoq	归属母公司股东的净利润环比增长率(%)(单季度) # equity_yoy	净资产同比增长率
    """
        Retrieves the financial data for a given stock within a specified date range.

        Args:
            stock_name (str): The stock code.
            start_date (str): The start date of the data range in the format "YYYYMMDD".
            end_date (str): The end date of the data range in the format "YYYYMMDD".
            financial_index (str, optional): The financial indicator to be queried.

        Returns:
            pd.DataFrame: A DataFrame containin financial data for the specified stock and date range.

"""
    stock_code = get_stock_code(stock_name)
    stock_data = pro.fina_indicator(**{
        "ts_code": stock_code,
        "ann_date": "",
        "start_date": start_date,
        "end_date": end_date,
        "period": '',
        "update_flag": "1",
        "limit": "",
        "offset": ""
    }, fields=["ts_code", "end_date", financial_index])

    #stock_name = get_stock_name_from_code(stock_code)
    stock_data['stock_name'] = stock_name
    stock_data = stock_data.sort_values(by='end_date', ascending=True)  # 按照日期升序排列
    # 把end_data列改名为trade_date
    stock_data.rename(columns={'end_date': 'trade_date'}, inplace=True)
    stock_financial_data = stock_data[['stock_name', 'trade_date', financial_index]]
    return stock_financial_data


def get_GDP_data(start_quarter:str='', end_quarter:str='', index:str='gdp_yoy') -> pd.DataFrame:
    # The available indicators for query include the following 9 categories: # gdp GDP累计值(亿元)# gdp_yoy 当季同比增速(%)# pi 第一产业累计值(亿元)# pi_yoy 第一产业同比增速(%)# si 第二产业累计值(亿元)# si_yoy 第二产业同比增速(%)# ti 第三产业累计值(亿元) # ti_yoy 第三产业同比增速(%)
    """
        Retrieves GDP data for the chosen index and specified time period.

        Args:
        - start_quarter (str): The start quarter of the query, in YYYYMMDD format.
        - end_quarter (str): The end quarter, in YYYYMMDD format.
        - index (str): The specific GDP index to retrieve. Default is `gdp_yoy`.

        Returns:
        - pd.DataFrame: A pandas DataFrame with three columns: `quarter`, `country`, and the selected `index`.
        """

    # The output is a DataFrame with three columns:
    # the first column represents the quarter (quarter), the second column represents the country (country), and the third column represents the index (index).
    df = pro.cn_gdp(**{
        "q":'',
        "start_q": start_quarter,
        "end_q": end_quarter,
        "limit": "",
        "offset": ""
    }, fields=[
        "quarter",
        "gdp",
        "gdp_yoy",
        "pi",
        "pi_yoy",
        "si",
        "si_yoy",
        "ti",
        "ti_yoy"
    ])
    df = df.sort_values(by='quarter', ascending=True)  #
    df['country'] = 'China'
    df = df[['quarter', 'country', index]].copy()


    return df

def get_cpi_ppi_currency_supply_data(start_month: str = '', end_month: str = '', type: str = 'cpi', index: str = '') -> pd.DataFrame:
    # The query types (type) include three categories: CPI, PPI, and currency supply. Each type corresponds to different indices.
    # Specifically, CPI has 12 indices, PPI has 30 indices, and currency supply has 9 indices.
    # The output is a DataFrame table with three columns: the first column represents the month (month), the second column represents the country (country), and the third column represents the index (index).

    # type='cpi',monthly CPI data include the following 12 categories:
    # nt_val	全国当月值 # nt_yoy	全国同比(%)# nt_mom	全国环比(%)# nt_accu	全国累计值# town_val	城市当月值# town_yoy	城市同比(%)# town_mom	城市环比(%)# town_accu	城市累计值# cnt_val	农村当月值# cnt_yoy	农村同比(%)# cnt_mom	农村环比(%)# cnt_accu	农村累计值

    # type = 'ppi', monthly PPI data include the following 30 categories:
    # ppi_yoy	PPI:全部工业品:当月同比
    # ppi_mp_yoy    PPI:生产资料:当月同比
    # ppi_mp_qm_yoy	PPI:生产资料:采掘业:当月同比
    # ppi_mp_rm_yoy	PPI:生产资料:原料业:当月同比
    # ppi_mp_p_yoy	PPI:生产资料:加工业:当月同比
    # ppi_cg_yoy	PPI:生活资料:当月同比
    # ppi_cg_f_yoy	PPI:生活资料:食品类:当月同比
    # ppi_cg_c_yoy	PPI:生活资料:衣着类:当月同比
    # ppi_cg_adu_yoy	PPI:生活资料:一般日用品类:当月同比
    # ppi_cg_dcg_yoy	PPI:生活资料:耐用消费品类:当月同比
    # ppi_mom	PPI:全部工业品:环比
    # ppi_mp_mom	PPI:生产资料:环比
    # ppi_mp_qm_mom	PPI:生产资料:采掘业:环比
    # ppi_mp_rm_mom	PPI:生产资料:原料业:环比
    # ppi_mp_p_mom	PPI:生产资料:加工业:环比
    # ppi_cg_mom	PPI:生活资料:环比
    # ppi_cg_f_mom	PPI:生活资料:食品类:环比
    # ppi_cg_c_mom	PPI:生活资料:衣着类:环比
    # ppi_cg_adu_mom	PPI:生活资料:一般日用品类:环比
    # ppi_cg_dcg_mom		PPI:生活资料:耐用消费品类:环比
    # ppi_accu		PPI:全部工业品:累计同比
    # ppi_mp_accu		PPI:生产资料:累计同比
    # ppi_mp_qm_accu		PPI:生产资料:采掘业:累计同比
    # ppi_mp_rm_accu		PPI:生产资料:原料业:累计同比
    # ppi_mp_p_accu	    PPI:生产资料:加工业:累计同比
    # ppi_cg_accu	PPI:生活资料:累计同比
    # ppi_cg_f_accu		PPI:生活资料:食品类:累计同比
    # ppi_cg_c_accu		PPI:生活资料:衣着类:累计同比
    # ppi_cg_adu_accu	PPI:生活资料:一般日用品类:累计同比
    # ppi_cg_dcg_accu	PPI:生活资料:耐用消费品类:累计同比

    # type = 'currency_supply', monthly currency supply data include the following 9 categories:
    # m0  M0(亿元)# m0_yoy  M0同比(%)# m0_mom  M0环比(%)# m1  M1(亿元)# m1_yoy  M1同比(%)# m1_mom  M1环比(%)# m2  M2(亿元)# m2_yoy  M2同比(%)# m2_mom  M2环比(%)

    """
        This function is used to retrieve China's monthly CPI (Consumer Price Index), PPI (Producer Price Index),
        and monetary supply data published by the National Bureau of Statistics,
        and return a DataFrame table containing month, country, and index values.
        The function parameters include start month, end month, query type, and query index.
        For query indexes that are not within the query range, the default index for the corresponding type is returned.

        Args:
        - start_month (str): start month of the query, in the format of YYYYMMDD.
        - end_month (str):end month in YYYYMMDD
        - type (str): required parameter, query type, including three types: cpi, ppi, and currency_supply.
        - index (str): optional parameter, query index, the specific index depends on the query type.
        If the query index is not within the range, the default index for the corresponding type is returned.

        Returns:
        - pd.DataFrame: DataFrame type, including three columns: month, country, and index value.
        """

    if type == 'cpi':

        df = pro.cn_cpi(**{
            "m": '',
            "start_m": start_month,
            "end_m": end_month,
            "limit": "",
            "offset": ""
        }, fields=[
            "month", "nt_val","nt_yoy", "nt_mom","nt_accu", "town_val", "town_yoy",  "town_mom",
            "town_accu", "cnt_val", "cnt_yoy", "cnt_mom", "cnt_accu"])
        # If the index is not within the aforementioned range, the index is set as "nt_yoy".
        if index not in df.columns:
            index = 'nt_yoy'


    elif type == 'ppi':
        df = pro.cn_ppi(**{
            "m": '',
            "start_m": start_month,
            "end_m": end_month,
            "limit": "",
            "offset": ""
        }, fields=[
            "month", "ppi_yoy", "ppi_mp_yoy", "ppi_mp_qm_yoy", "ppi_mp_rm_yoy", "ppi_mp_p_yoy", "ppi_cg_yoy",
            "ppi_cg_f_yoy", "ppi_cg_c_yoy", "ppi_cg_adu_yoy", "ppi_cg_dcg_yoy",
            "ppi_mom", "ppi_mp_mom", "ppi_mp_qm_mom", "ppi_mp_rm_mom", "ppi_mp_p_mom", "ppi_cg_mom", "ppi_cg_f_mom",
            "ppi_cg_c_mom", "ppi_cg_adu_mom", "ppi_cg_dcg_mom",
            "ppi_accu", "ppi_mp_accu", "ppi_mp_qm_accu", "ppi_mp_rm_accu", "ppi_mp_p_accu", "ppi_cg_accu",
            "ppi_cg_f_accu", "ppi_cg_c_accu", "ppi_cg_adu_accu", "ppi_cg_dcg_accu"
        ])
        if index not in df.columns:
            index = 'ppi_yoy'

    elif type == 'currency_supply':
        df = pro.cn_m(**{
            "m": '',
            "start_m": start_month,
            "end_m": end_month,
            "limit": "",
            "offset": ""
        }, fields=[
            "month", "m0",  "m0_yoy","m0_mom", "m1",
            "m1_yoy",  "m1_mom", "m2", "m2_yoy", "m2_mom"])
        if index not in df.columns:
            index = 'm2_yoy'


    df = df.sort_values(by='month', ascending=True)  #
    df['country'] = 'China'
    df = df[['month', 'country', index]].copy()
    return df

def predict_next_value(df: pd.DataFrame, pred_index: str = 'nt_yoy', pred_num:int = 1. ) -> pd.DataFrame:
    """
    Predict the next n values of a specific column in the DataFrame using linear regression.

        Parameters:
            df (pandas.DataFrame): The input DataFrame.
            pred_index (str): The name of the column to predict.
            pred_num (int): The number of future values to predict.

        Returns:
        pandas.DataFrame: The DataFrame with the predicted values appended to the specified column
                          and other columns filled as pred+index.
        """
    input_array = df[pred_index].values

    # Convert the input array into the desired format.
    x = np.array(range(len(input_array))).reshape(-1, 1)
    y = input_array.reshape(-1, 1)

    # Train a linear regression model.
    model = LinearRegression()
    model.fit(x, y)

    # Predict the future n values.
    next_indices = np.array(range(len(input_array), len(input_array) + pred_num)).reshape(-1, 1)
    predicted_values = model.predict(next_indices).flatten()

    for i, value in enumerate(predicted_values, 1):
        row_data = {pred_index: value}
        for other_col in df.columns:
            if other_col != pred_index:
                row_data[other_col] = 'pred' + str(i)
        df = df.append(row_data, ignore_index=True)

        # Return the updated DataFrame
    return df






def get_latest_new_from_web(src: str = 'sina') -> pd.DataFrame:

    # 新浪财经	sina	获取新浪财经实时资讯
    # 同花顺	    10jqka	同花顺财经新闻
    # 东方财富	eastmoney	东方财富财经新闻
    # 云财经	    yuncaijing	云财经新闻
    """
    Retrieves the latest news data from major news websites, including Sina Finance, 10jqka, Eastmoney, and Yuncaijing.

    Args:
        src (str): The name of the news website. Default is 'sina'. Optional parameters include: 'sina' for Sina Finance,
        '10jqka' for 10jqka, 'eastmoney' for Eastmoney, and 'yuncaijing' for Yuncaijing.

    Returns:
        pd.DataFrame: A DataFrame containing the news data, including two columns for date/time and content.
    """

    df = pro.news(**{
        "start_date": '',
        "end_date": '',
        "src": src,
        "limit": "",
        "offset": ""
    }, fields=[
        "datetime",
        "content",
    ])
    df = df.apply(lambda x: '[' + x.name + ']' + ': ' + x.astype(str))
    return df


# def show_dynamic_table(df: pd.DataFrame) -> None:
#     '''
#      This function displays a dynamic table in the terminal window, where each row of the input DataFrame is shown one by one.
#      Arguments:
#         df: A Pandas DataFrame containing the data to be displayed in the dynamic table.
#
#      Returns: None. This function does not return anything.
#
#     '''
#
#     return df
#     # table = PrettyTable(df.columns.tolist(),align='l')
#
#     # 将 DataFrame 的数据添加到表格中
#     # for row in df.itertuples(index=False):
#     #     table.add_row(row)
#
#     # 初始化终端
#     # term = Terminal()
#     #
#     # # 在终端窗口中滚动显示表格
#     # with term.fullscreen():
#     #     with term.cbreak():
#     #         print(term.clear())
#     #         with term.location(0, 0):
#     #             # 将表格分解为多行,并遍历每一行
#     #             lines = str(table).split('\n')
#     #             for i, line in enumerate(lines):
#     #                 with term.location(0, i):
#     #                     print(line)
#     #                     time.sleep(1)
#     #
#     #             while True:
#     #                 # 读取输入
#     #                 key = term.inkey(timeout=0.1)
#     #
#     #                 # 如果收到q键,则退出
#     #                 if key.lower() == 'q':
#     #                     break


def get_index_constituent(index_name: str = '', start_date:str ='', end_date:str ='') -> pd.DataFrame:
    """
        Query the constituent stocks of basic index (中证500) or a specified SW (申万) industry index

        args:
             index_name: the name of the index.
             start_date: the start date in "YYYYMMDD".
             end_date:  the end date in "YYYYMMDD".

        return:
            A pandas DataFrame containing the following columns:
            index_code
            index_name
            stock_code: the code of the constituent stock.
            stock_name:  the name of the constituent stock.
            weight: the weight of the constituent stock.
    """

    if '申万' in index_name:
        if '申万一级行业' in index_name:
            # index_name取后面的名字
            index_name = index_name[6:]
            df1 = pd.read_csv('SW2021_industry_L1.csv')
            index_code = df1[df1['industry_name'] == index_name]['index_code'].iloc[0]
        elif '申万二级行业' in index_name:
            index_name = index_name[6:]
            df1 = pd.read_csv('SW2021_industry_L2.csv')
            index_code = df1[df1['industry_name'] == index_name]['index_code'].iloc[0]
        elif '申万三级行业' in index_name:
            index_name = index_name[6:]
            df1 = pd.read_csv('SW2021_industry_L3.csv')
            index_code = df1[df1['industry_name'] == index_name]['index_code'].iloc[0]

        print('The industry code for ', index_name, ' is: ', index_code)

        # 拉取数据
        df = pro.index_member(**{
            "index_code": index_code ,  #'851251.SI'
            "is_new": "",
            "ts_code": "",
            "limit": "",
            "offset": ""
        }, fields=[
            "index_code",
            "con_code",
            "in_date",
            "out_date",
            "is_new",
            "index_name",
            "con_name"
        ])
        #
        # For each stock, filter the start_date and end_date that are between in_date and out_date.
        df = df[(df['in_date'] <= start_date)]
        df = df[(df['out_date'] >= end_date) | (df['out_date'].isnull())]



        df.rename(columns={'con_code': 'stock_code'}, inplace=True)

        df.rename(columns={'con_name': 'stock_name'}, inplace=True)
        #
        df['weight'] = np.nan

        df = df[['index_code', "index_name", 'stock_code', 'stock_name','weight']]

    else: # 宽基指数
        df1 = pro.index_basic(**{
            "ts_code": "",
            "market": "",
            "publisher": "",
            "category": "",
            "name": index_name,
            "limit": "",
            "offset": ""
        }, fields=[
            "ts_code",
            "name",
        ])

        index_code = df1["ts_code"][0]
        print(f'index_code for basic index {index_name} is {index_code}')


        # Step 2: Retrieve the constituents of an index based on the index code and given date.
        df = pro.index_weight(**{
            "index_code": index_code,
            "trade_date": '',
            "start_date": start_date,
            "end_date": end_date,
            "limit": "",
            "offset": ""
        }, fields=[
            "index_code",
            "con_code",
            "trade_date",
            "weight"
        ])
        # df = df.sort_values(by='trade_date', ascending=True)  #
        df['index_name'] = index_name
        last_day = df['trade_date'][0]
        #  for the last trading day
        df = df[df['trade_date'] == last_day]
        df_stock = pd.read_csv('tushare_stock_basic_20230421210721.csv')
        # Merge based on the stock code.
        df = pd.merge(df, df_stock, how='left', left_on='con_code', right_on='ts_code')
        # df.rename(columns={'name_y': 'name'}, inplace=True)
        df = df.drop(columns=['symbol', 'area', 'con_code'])
        df.sort_values(by='weight', ascending=False, inplace=True)
        df.rename(columns={'name': 'stock_name'}, inplace=True)
        df.rename(columns={'ts_code': 'stock_code'}, inplace=True)
        df.dropna(axis=0, how='any', inplace=True)
        #
        df = df[['index_code', "index_name", 'stock_code', 'stock_name', 'weight']]

    return df

# Determine whether the given name is a stock or a fund.,
def is_fund(ts_name: str = '') -> bool:
    # call  get_stock_code()和query_fund_name_or_code()
    if get_stock_code(ts_name) is not None and query_fund_name_or_code(ts_name) is None:
        return False
    elif get_stock_code(ts_name) is None and query_fund_name_or_code(ts_name) is not None:
        return True




def calculate_earning_between_two_time(stock_name: str = '', start_date: str = '', end_date: str = '', index: str = 'close') -> float:
    """
        Calculates the rate of return for a specified stock/fund between two dates.

        Args:
            stock_name: stock_name or fund_name
            start_date
            end_date
            index (str): The index used to calculate the stock return, including 'open' and 'close'.

        Returns:
            float: The rate of return for the specified stock between the two dates.
    """
    if is_fund(stock_name):
        fund_code = query_fund_name_or_code(stock_name)
        stock_data = query_fund_data(fund_code, start_date, end_date)
        if index =='':
            index = 'adj_nav'
    else:
        stock_data = get_stock_prices_data(stock_name, start_date, end_date,'daily')
    try:
        end_price = stock_data.iloc[-1][index]
        start_price = stock_data.iloc[0][index]
        earning = cal_dt(end_price, start_price)
        # earning = round((end_price - start_price) / start_price * 100, 2)
    except:
        print(ts_code,start_date,end_date)
        print('##################### 该股票没有数据 #####################')
        return None
    # percent = earning * 100
    # percent_str = '{:.2f}%'.format(percent)

    return  earning


def loop_rank(df: pd.DataFrame,  func: callable, *args, **kwargs) -> pd.DataFrame:
    """
        It iteratively applies the given function to each row and get a result using function. It then stores the calculated result in 'new_feature' column.

        Args:
        df: DataFrame with a single column
        func : The function to be applied to each row: func(row, *args, **kwargs)
        *args: Additional positional arguments for `func` function.
        **kwargs: Additional keyword arguments for `func` function.

        Returns:
        pd.DataFrame: A output DataFrame with three columns: the constant column, input column, and new_feature column.
                     The DataFrame is sorted based on the new_feature column in descending order.

        """
    df['new_feature'] = None
    loop_var = df.columns[0]
    for _, row in df.iterrows():
        res  = None
        var = row[loop_var]                                         #

        if var is not None:
            if loop_var == 'stock_name':
                stock_name = var
            elif loop_var == 'stock_code':
                stock_name = get_stock_name_from_code(var)
            elif loop_var == 'fund_name':
                stock_name = var
            elif loop_var == 'fund_code':
                stock_name = query_fund_name_or_code('',var)
            time.sleep(0.4)
            try:
                res = func(stock_name, *args, **kwargs)             #
            except:
                raise ValueError('#####################Error for func#####################')
            # res represents the result obtained for the variable. For example, if the variable is a stock name, res could be the return rate of that stock over a certain period or a specific feature value of that stock. Therefore, res should be a continuous value.
            # If the format of res is a float, then it can be used directly. However, if res is in DataFrame format, you can retrieve the value corresponding to the index.
            if isinstance(res, pd.DataFrame) and not res.empty:
                #
                try:
                    res = round(res.loc[:,args[-1]][0], 2)
                    df.loc[df[loop_var] == var, 'new_feature'] = res
                except:
                    raise ValueError('##################### Error ######################')
            elif isinstance(res, float): #
                res = res
                df.loc[df[loop_var] == var, 'new_feature'] = res
            print(var, res)


    # Remove the rows where the new_feature column is empty.
    df = df.dropna(subset=['new_feature'])
    stock_data = df.sort_values(by='new_feature', ascending=False)
    #
    stock_data.insert(0, 'unchanged', loop_var)
    stock_data = stock_data.loc[:,[stock_data.columns[0], loop_var, 'new_feature']]

    return stock_data

def output_mean_median_col(data: pd.DataFrame, col: str = 'new_feature') -> float:
    # It calculates the mean and median value for the specified column.

    mean = round(data[col].mean(), 2)
    median = round(data[col].median(), 2)
    #
    #print(title, mean)
    return (mean, median)


# def output_median_col(data: pd.DataFrame, col: str, title_name: str = '') -> float:
#     # It calculates the median value for the specified column and returns the median as a float value.
#
#     median = round(data[col].median(), 2)
#     #print(title_name, median)
#
#     return median


def output_weighted_mean_col(data: pd.DataFrame, col: str, weight_col: pd.Series) -> float:

    """
        Calculates the weighted mean of a column and returns the result as a float.

        Args:
            data (pd.DataFrame): The input cross-sectional or time-series data containing the feature columns.
            col (str): The name of the feature column to calculate the weighted mean for.
            weight_col (pd.Series): The weights used for the calculation, as a pandas Series.

        Returns:
            float: The weighted mean of the specified feature column.
        """

    weighted_mean = round(np.average(data[col], weights = weight_col)/100., 2)
    return weighted_mean



def get_index_data(index_name: str = '', start_date: str = '', end_date: str = '', freq: str = 'daily') -> pd.DataFrame:
    """
        This function retrieves daily, weekly, or monthly data for a given stock index.

        Arguments:
        - index_name: Name of the index
        - start_date: Start date in 'YYYYMMDD'
        - end_date: End date in 'YYYYMMDD'
        - freq: Frequency 'daily', 'weekly', or 'monthly'

        Returns:
        A DataFrame containing the following columns:
        trade_date, ts_code, close, open, high, low, pre_close: Previous day's closing price, change(涨跌额), pct_chg(涨跌幅), vol(成交量), amount(成交额), name: Index Name
        """
    df1 = pro.index_basic(**{
        "ts_code": "",
        "market": "",
        "publisher": "",
        "category": "",
        "name": index_name,
        "limit": "",
        "offset": ""
    }, fields=[
        "ts_code",
        "name",
    ])

    index_code = df1["ts_code"][0]
    print(f'index_code for index {index_name} is {index_code}')
    #
    if freq == 'daily':
        df = pro.index_daily(**{
            "ts_code": index_code,
            "trade_date": '',
            "start_date": start_date,
            "end_date": end_date,
            "limit": "",
            "offset": ""
        }, fields=[
            "trade_date",
            "ts_code",
            "close",
            "open",
            "high",
            "low",
            "pre_close",
            "change",
            "pct_chg",
            "vol",
            "amount"
        ])
    elif freq == 'weekly':
        df = pro.index_weekly(**{
            "ts_code": index_code,
            "trade_date": '',
            "start_date": start_date,
            "end_date": end_date,
            "limit": "",
            "offset": ""
        }, fields=[
            "trade_date",
            "ts_code",
            "close",
            "open",
            "high",
            "low",
            "pre_close",
            "change",
            "pct_chg",
            "vol",
            "amount"
        ])
    elif freq == 'monthly':
        df = pro.index_monthly(**{
            "ts_code": index_code,
            "trade_date": '',
            "start_date": start_date,
            "end_date": end_date,
            "limit": "",
            "offset": ""
        }, fields=[
            "trade_date",
            "ts_code",
            "close",
            "open",
            "high",
            "low",
            "pre_close",
            "change",
            "pct_chg",
            "vol",
            "amount"
        ])

    df = df.sort_values(by='trade_date', ascending=True)  #
    df['index_name'] = index_name
    return df





def get_north_south_money(start_date: str = '', end_date: str = '', trade_date: str = '') -> pd.DataFrame:
    #
    # trade_date: 交易日期
    # ggt_ss:	港股通(上海)
    # ggt_sz:	港股通(深圳)
    # hgt:	沪股通(亿元)
    # sgt:	深股通(亿元)
    # north_money:	北向资金(亿元)= hgt + sgt
    # south_money:	南向资金(亿元)= ggt_ss + ggt_sz
    # name:  固定为'A-H',代表A股和H股
    # accumulate_north_money: 累计北向资金流入
    # accumulate_south_money: 累计南向资金流入


    month_df = pro.moneyflow_hsgt(**{
        "trade_date": trade_date,
        "start_date": start_date,
        "end_date": end_date,
        "limit": "",
        "offset": ""
    }, fields=[
        "trade_date",
        "ggt_ss",
        "ggt_sz",
        "hgt",
        "sgt",
        "north_money",
        "south_money"
    ])

    month_df[['ggt_ss','ggt_sz','hgt','sgt','north_money','south_money']] = month_df[['ggt_ss','ggt_sz','hgt','sgt','north_money','south_money']]/100.0
    month_df = month_df.sort_values(by='trade_date', ascending=True)  #
    month_df['stock_name'] = 'A-H'
    month_df['accumulate_north_money'] = month_df['north_money'].cumsum()
    month_df['accumulate_south_money'] = month_df['south_money'].cumsum()
    return month_df



def plot_k_line(stock_data: pd.DataFrame, title: str = '') -> None:
    """
        Plots a K-line chart of stock price and volume.

        Args:
            stock_data : A pandas DataFrame containing the stock price information, in which each row
                represents a daily record. The DataFrame must contain the 'trade_date','open', 'close', 'high', 'low','volume', 'name' columns, which is used for k-line and volume.
                如果dataframe中还含有'macd','kdj', 'rsi', 'cci', 'boll','pe_ttm','turnover_rate'等列,则在k线图下方绘制这些指标的子图.
            title : The title of the K-line chart.

        Returns:
            None
    """

    #
    stock_data['trade_date'] = pd.to_datetime(stock_data['trade_date'], format='%Y%m%d')
    stock_data.set_index('trade_date', inplace=True)
    #
    custom_style = mpf.make_marketcolors(up='r', down='k', inherit=True)
    china_style = mpf.make_mpf_style(marketcolors=custom_style)

    # MACD
    # stock_data['macd1'] = stock_data['Close'].ewm(span=12).mean() - stock_data['Close'].ewm(span=26).mean()
    # stock_data['macd_signal1'] = stock_data['macd'].ewm(span=9).mean()

    #
    #mpf.plot(stock_data, type='candle', volume=True, title=title, mav=(5, 10, 20), style = china_style, addplot = macd)
    add_plot = []
    # The index column is located after the name column in the last few columns.
    # Retrieve the column names after the 'name' column.
    index_list = stock_data.columns[stock_data.columns.get_loc('stock_name')+1:]

    index_df = stock_data[index_list]

    color_list = ['green','blue','red','yellow','black','purple','orange','pink','brown','gray']
    custom_lines = []
    for i in range(len(index_list)):
        # If the column names contain 'boll', set panel to 0. Otherwise, set panel to 2.
        if 'boll' in index_list[i]:
            sub_plot = mpf.make_addplot(index_df[index_list[i]], panel=0, ylabel=index_list[i], color=color_list[i], type='line', secondary_y=True)
        elif  index_list[i] =='macd':
            sub_plot = mpf.make_addplot(index_df[index_list[i]], panel=2, ylabel=index_list[i], color=color_list[i], type='bar', secondary_y=False)

        else:
            sub_plot = mpf.make_addplot(index_df[index_list[i]], panel=2, ylabel=index_list[i], color=color_list[i], type='line', secondary_y=False)

        custom_line = Line2D([0], [0], color=color_list[i], lw=1, linestyle='dashed')


        add_plot.append(sub_plot)
        custom_lines.append(custom_line)

    mav_colors = ['red', 'green', 'blue']

    fig, axes = mpf.plot(stock_data, type='candle', volume=True, title=title, mav=(5, 10, 20), mavcolors=mav_colors, style=china_style, addplot=add_plot, returnfig=True)


    mav_labels = ['5-day MA', '10-day MA', '20-day MA']
    #
    legend_lines = [plt.Line2D([0], [0], color=color, lw=2) for color in mav_colors]

    #
    axes[0].legend(legend_lines, mav_labels)

    if len(index_list) ==1:
        label = index_list[0]
    elif len(index_list) > 1:
        label_list = [i.split('_')[0] for i in index_list]
        #
        label = list(set(label_list))[0]

    if len(index_list) >= 1:
        if 'boll' in label:
            axes[0].legend(custom_lines, index_list, loc='lower right')

        elif len(index_list) > 1:
            axes[-2].set_ylabel(label)
            axes[-2].legend(custom_lines, index_list,  loc='lower right')

    #
    fig.set_size_inches(20, 16)
    #
    for ax in axes:
        ax.grid(True)

    #fig.show()
    return axes


def cal_dt(num_at_time_2: float = 0.0, num_at_time_1: float = 0.0) -> float:
    """
        This function calculates the percentage change of a metric from one time to another.

        Args:
        - num_at_time_2: the metric value at time 2 (end time)
        - num_at_time_1: the metric value at time 1 (start time)

        Returns:
        - float: the percentage change of the metric from time 1 to time 2

        """
    if num_at_time_1 == 0:
        num_at_time_1 = 0.0000000001
    return round((num_at_time_2 - num_at_time_1) / num_at_time_1, 4)


def query_fund_info(fund_code: str = '') -> pd.DataFrame:
    #
    # fund_code	str	Y	基金代码 # fund_name	str	Y	简称 # management	str	Y	管理人 # custodian	str	Y	托管人 # fund_type	str	Y	投资类型 # found_date	str	Y	成立日期 # due_date	str	Y	到期日期 # list_date	str	Y	上市时间 # issue_date	str	Y	发行日期 # delist_date	str	Y	退市日期 # issue_amount	float	Y	发行份额(亿) # m_fee	float	Y	管理费 # c_fee	float	Y	托管费
    # duration_year	float	Y	存续期 # p_value	float	Y	面值 # min_amount	float	Y	起点金额(万元) # benchmark	str	Y	业绩比较基准 # status	str	Y	存续状态D摘牌 I发行 L已上市 # invest_type	str	Y	投资风格 # type	str	Y	基金类型 # purc_startdate	str	Y	日常申购起始日 # redm_startdate	str	Y	日常赎回起始日 # market	str	Y	E场内O场外
    """
        Retrieves information about a fund based on the fund code.

        Args:
            fund_code (str, optional): Fund code. Defaults to ''.

        Returns:
            df (DataFrame): A DataFrame containing various information about the fund, including fund code, fund name,
                            management company, custodian company, investment type, establishment date, maturity date,
                            listing date, issuance date, delisting date, issue amount, management fee, custodian fee,
                            fund duration, face value, minimum investment amount, benchmark, fund status, investment style,
                            fund type, start date for daily purchases, start date for daily redemptions, and market type.
                            The column 'ts_code' is renamed to 'fund_code', and 'name' is renamed to 'fund_name' in the DataFrame.
        """
    df = pro.fund_basic(**{
        "ts_code": fund_code,
        "market": "",
        "update_flag": "",
        "offset": "",
        "limit": "",
        "status": "",
        "name": ""
    }, fields=[
        "ts_code",
        "name",
        "management",
        "custodian",
        "fund_type",
        "found_date",
        "due_date",
        "list_date",
        "issue_date",
        "delist_date",
        "issue_amount",
        "m_fee",
        "c_fee",
        "duration_year",
        "p_value",
        "min_amount",
        "benchmark",
        "status",
        "invest_type",
        "type",
        "purc_startdate",
        "redm_startdate",
        "market"
    ])
    #
    df.rename(columns={'ts_code': 'fund_code'}, inplace=True)
    df.rename(columns={'name': 'fund_name'}, inplace=True)
    return df

def query_fund_data(fund_code: str = '', start_date: str = '', end_date: str = '') -> pd.DataFrame:
    #
    # ts_code	str	Y	TS代码 # ann_date	str	Y	公告日期 # nav_date	str	Y	净值日期 # unit_nav	float	Y	单位净值 # accum_nav	float	Y	累计净值
    # accum_div	float	Y	累计分红 # net_asset	float	Y	资产净值 # total_netasset	float	Y	合计资产净值 # adj_nav	float	Y	复权单位净值  pct_chg 每日涨跌幅
    """
        Retrieves fund data based on the fund code, start date, and end date.

        Args:
            fund_code (str, optional): Fund code. Defaults to ''.
            start_date (str, optional): Start date in YYYYMMDD format. Defaults to ''.
            end_date (str, optional): End date in YYYYMMDD format. Defaults to ''.

        Returns:
            df (DataFrame): A DataFrame containing fund data, including TS code, announcement date, net asset value date,
                            unit net asset value, accumulated net asset value, accumulated dividends, net asset value,
                            total net asset value, adjusted unit net asset value, and fund name. The 'ts_code' column is renamed
                            to 'fund_code', 'nav_date' is renamed to 'trade_date', and the DataFrame is sorted by the trade date
                            in ascending order. If the fund code does not exist, None is returned.
        """
    df = pro.fund_nav(**{
        "ts_code": fund_code,
        "nav_date": "",
        "offset": "",
        "limit": "",
        "market": "",
        "start_date": start_date,
        "end_date": end_date
    }, fields=[
        "ts_code",
        "ann_date",
        "nav_date",
        "unit_nav",
        "accum_nav",
        "accum_div",
        "net_asset",
        "total_netasset",
        "adj_nav",
        "update_flag"
    ])
    try:
        fund_name= query_fund_name_or_code(fund_code=fund_code)
        df['fund_name'] = fund_name
        #
        df.rename(columns={'ts_code': 'fund_code'}, inplace=True)
        df.rename(columns={'nav_date': 'trade_date'}, inplace=True)
        df.sort_values(by='trade_date', ascending=True, inplace=True)
    except:
        print(fund_code,'基金代码不存在')
        return None
    #
    df['pct_chg'] = df['adj_nav'].pct_change()
    #
    df.loc[0, 'pct_chg'] = 0.0


    return df

def query_fund_name_or_code(fund_name: str = '', fund_code: str = '') -> str:
    #
    """
        Retrieves the fund code based on the fund name or Retrieves the fund name based on the fund code.

        Args:
            fund_name (str, optional): Fund name. Defaults to ''.
            fund_code (str, optional): Fund code. Defaults to ''.

        Returns:
            code or name: Fund code if fund_name is provided and fund_code is empty. Fund name if fund_code is provided and fund_name is empty.
        """


    #df = pd.read_csv('./tushare_fund_basic_20230508193747.csv')
    # Query the fund code based on the fund name.
    if fund_name != '' and fund_code == '':
        #
        df = pd.read_csv('./tushare_fund_basic_all.csv')
        #
        # df = pro.fund_basic(**{
        #     "ts_code": "",
        #     "market": "",
        #     "update_flag": "",
        #     "offset": "",
        #     "limit": "",
        #     "status": "",
        #     "name": fund_name
        # }, fields=[
        #     "ts_code",
        #     "name"
        # ])
        try:
            #
            code = df[df['name'] == fund_name]['ts_code'].values[0]
        except:
            #print(fund_name,'基金名称不存在')
            return None
        return code
    # Query the fund name based on the fund code.
    if fund_code != '' and fund_name == '':
        df = pd.read_csv('./tushare_fund_basic_all.csv')
        try:
            name = df[df['ts_code'] == fund_code]['name'].values[0]
        except:
            #print(fund_code,'基金代码不存在')
            return None
        return name



def print_save_table(df: pd.DataFrame, title_name: str, save:bool = False ,file_path: str = './output/') -> None:
    """
        It prints the dataframe as a formatted table using the PrettyTable library and saves it to a CSV file at the specified file path.

        Args:
        - df: the dataframe to be printed and saved to a CSV file
        - title_name: the name of the table to be printed and saved
        - save: whether to save the table to a CSV file
        - file_path:  the file path where the CSV file should be saved.

        Returns: None
    """

    # 创建表格table.max_width = 20

    # table = PrettyTable(df.columns.tolist())
    # table.align = 'l'
    # table.max_width = 40
    #
    # #
    # for row in df.itertuples(index=False):
    #     table.add_row(row)

    #print(table)


    if not os.path.exists(file_path):
        os.makedirs(file_path)

    if file_path is not None and save == True:
        file_path = file_path + title_name + '.csv'
        df.to_csv(file_path, index=False)
    return df



#
def merge_indicator_for_same_stock(df1: pd.DataFrame, df2: pd.DataFrame) -> pd.DataFrame:
    """
        Merges two DataFrames (two indicators of the same stock) based on common names for same stock. Data from two different stocks cannot be merged

        Args:
            df1: DataFrame contains some indicators for stock A.
            df2: DataFrame contains other indicators for stock A.

        Returns:
            pd.DataFrame: The merged DataFrame contains two different indicators.
    """
    if len(set(df1.columns).intersection(set(df2.columns))) > 0:
        # If there are identical column names, merge the two DataFrames based on the matching column names.
        #
        common_cols = list(set(df1.columns).intersection(set(df2.columns)))
        #
        df = pd.merge(df1, df2, on=common_cols)
        return  df
    else:
        #
        raise ValueError('The two dataframes have no columns in common.')

def select_value_by_column(df1:pd.DataFrame, col_name: str = '', row_index: int = -1) -> Union[pd.DataFrame, Any]:
    """
        Selects a specific column or a specific value within a DataFrame.

        Args:
            df1: The input DataFrame.
            col_name: The name of the column to be selected.
            row_index: The index of the row to be selected.

        Returns:
            Union[pd.DataFrame, Any]. row_index=-1: df1[col_name].to_frame() or df1[col_name][row_index]
    """
    if row_index == -1:
        #
        return df1[col_name].to_frame()
    else:
        #
        return df1[col_name][row_index]



if __name__ == "__main__":
    stock_name='成都银行'
    stock_name2='五粮液'
    stock_name3 = '宁德时代'
    start = '20230104'
    end = '20230504'
    fund_name = "华商优势行业" #'易方达蓝筹精选'

    start_quarter = '201001'
    end_quarter = '202303'
    title_name ='上证50成分股收益率'
    ax = None
    res = is_fund('易方达蓝筹精选')
    #_, ax = plt.subplots()
    # code = query_fund_name_or_code('华商优势行业')
    # ------------step1 数据查询层  获取股票代码
    # start_last_year = get_last_year_date(start)
    # end_last_year = get_last_year_date(end)
    stock_code = get_stock_code(stock_name)
    # name = get_stock_name_from_code(stock_code)
    # print(name)
    # print(stock_code)
    # stock_code2 = get_stock_code(stock_name2)
    # stock_code3 = get_stock_code(stock_name3)
    # stock_technical_data = get_Financial_data(stock_code, start, end)
    # macrodata = get_ppi_data('', start_quarter, end_quarter, 'ppi_yoy')
    # index_daily = get_index_data('沪深300',start,end,'daily')
    # index_daily2 = get_index_data('中证500',start,end,'daily')
    # index_daily3 = get_index_data('中证1000',start,end,'daily')
    # index_daily4 = get_index_data('创业板指',start,end,'daily')
    #stock_data = get_index_constituent('上证50','20230101','20230508')
    # money = get_north_south_money('20230425', '20230426')
    # stock_data = get_stock_prices_data(stock_code, start, end)

    # stock_data = get_stock_monthly_prices_data("","", "",'20230331')
    # stock_data = get_stock_prices_data('', start, end, 'daily')
    # fund_df = query_fund_Manager('周海栋')
    #
    # fund_code  = select_value_by_column(fund_df, 'fund_code', -1)
    # res_earning = loop_rank(fund_code, calculate_earning_between_two_time, start, end, 'adj_nav')
    # print(res_earning)
    #fund_code = query_fund_name_or_code(fund_name,'')


    #fund_data = query_fund_data(fund_code, start, end)
    #df_daily = get_daily_trading_data(stock_code,'20200101', '20230526')
    # stock_data2 = get_stock_prices_data(stock_code2, start, end,'daily')
    # stock_data3 = get_stock_prices_data(stock_code3, start, end,'daily')

    # dynamic_new = get_latest_new_from_web('sina')
    #stock_df = get_sw_industry_stock('城商行Ⅱ','L2')
    # df_macro = get_cpi_ppi_currency_supply_data('200101','202304','cpi','nt_yoy')
    # df_macro = get_cpi_ppi_currency_supply_data('200101','202304','ppi','ppi_yoy')
    # df_macro = get_cpi_ppi_currency_supply_data('200101','202304','currency_supply','m2_yoy')
    # df_gdp = get_GDP_data('2001Q1','2023Q1','gdp_yoy')
    # df_gdp = predict_next_value(df_gdp, 'gdp_yoy', 4)
    #company_df = get_company_info('贵州茅台')
    #print_save_table(company_df, '贵州茅台公司信息')
    #fin_df = get_Financial_data_from_time_range(stock_code, '20200101', '20230526','roe')

    #tech_df = get_stock_technical_data(stock_code, start, end)



    # ----------------------------------step2 数据处理层  在截面或者时序数据-------------------------------------------------------
    # 提取相应指标, 数据处理, 排序,提取,求差,加工..,
    # fund_info = query_fund_info('005827.OF')
    # value = select_value_by_column(fund_info, 'fund_name', 0)
    #fund_index = calculate_stock_index(fund_data,'adj_nav')
    #stock_index = rank_index_cross_section(stock_data, 'pct_chg', -1, False)
    #stock_index = calculate_stock_index(stock_data, 'pct_chg')
    #stock_index_each_day = calculate_stock_index(money, 'north_money')
    #stock_index = calculate_stock_index(fin_df, 'roe')
    # stock_index2 = calculate_stock_index(stock_data2, 'Cumulative_Earnings_Rate')
    # stock_index3 = calculate_stock_index(stock_data3, 'Cumulative_Earnings_Rate')
    # stock_index4 = calculate_stock_index(index_daily4, 'Cumulative_Earnings_Rate')
    # stock_index2 = calculate_stock_index(stock_data2, 'Cumulative_Earnings_Rate')
    #stock_index = calculate_stock_index(stock_data1, 'close')
    #stock_index2 = calculate_stock_index(tech_df, 'macd')
    #stock_index1 = calculate_stock_index(stock_data, 'candle_K')
    #stock_index2 = calculate_stock_index(df_daily, 'pe_ttm')
    #merge_df = merge_data(stock_index1, stock_index2)
    #res_earning = loop_rank(stock_data, 'stock_name', calculate_earning_between_two_time, start, end)
    # index_profit_yoy = loop_rank(stock_data, 'stock_name', get_Financial_data, start, end, 'profit_dedt')
    # index_profit_yoy = loop_rank(stock_data, 'stock_name', get_Financial_data, start, end, 'netprofit_yoy')

    #res_earning_top_n = rank_index_cross_section(stock_index, 10, False)
    #index_profit_yoy_last = loop_rank(stock_data, 'stock_name', get_Financial_data, start_last_year, end_last_year, 'profit_dedt')
    # profit_yoy = calculate_stock_index(stock_technical_data, 'dt_netprofit_yoy')
    # accumulate_north_month = calculate_stock_index(money, 'accumulate_south_money')
    # accumulate_north_month = calculate_stock_index(res_earning, 'accumulate_south_money')
    # stock_code = get_stock_code(stock_name)
    # fin_df1 = get_Financial_data_from_time_range(stock_code, '20150101', '20230526', 'roa')
    # fin_df2 = get_Financial_data_from_time_range(stock_code, '20150101', '20230526', 'roa')
    # ax = plot_stock_data(fin_df1, ax, 'line', title_name)
    # ax = plot_stock_data(fin_df2, ax, 'line', title_name)
    #stock_data = get_index_constituent('上证50','20220105', '20230505')
    # stock_data = get_index_constituent('申万二级行业城商行Ⅱ','20220105', '20220505')
    # #stock_list = select_value_by_column(stock_data, 'stock_name', -1)
    #
    # index_profit_yoy = loop_rank(stock_list, get_Financial_data, start, 'netprofit_yoy')
    # median = output_median_col(index_profit_yoy, 'new_feature')
    # ax = plot_stock_data(index_profit_yoy, ax, 'bar', '上证50的最近季度归母净利润同比增长率')






    # ----------------------------------step3 可视化层:文字,图片,表格等多种模态数据输出-------------------------------------------------------
    #ax = plot_stock_data(stock_index, ax, 'line', title_name)
    #ax = plot_stock_data(stock_index_each_day, ax, 'bar', title_name)
    #print_save_table(fund_info, title_name)

    #_, sum_new = output_mean_sum_col(index_profit_yoy,'new_feature')
    #_, sum_old = output_mean_sum_col(index_profit_yoy_last,'new_feature')


    #print('科创50成分股的最近季度归母净利润同比增长率中位数%:', median)
    #dt = cal_dt(sum_new, sum_old)
    #print('上证50成分股的最近季度归母净利润同比增长率:',dt)

    #plot_k_line(merge_df, title_name)
    # ax = plot_stock_data(index_profit_yoy, ax, 'bar', '上证50成分股的最近季度归母净利润同比增长率')
    #ax = plot_stock_data(accumulate_north_month, ax, 'line', '2023年1月至4月南向资金累计流向')

    # ax2 = plot_stock_data(stock_index2, ax1, 'line', '贵州茅台VS五粮液近十年收益率对比图')
    # ax = plot_stock_data(stock_index, ax,'line', title_name)
    # ax = plot_stock_data(stock_index2, ax,'line', title_name)
    # ax = plot_stock_data(stock_index3, ax,'line', title_name)
    # ax = plot_stock_data(stock_index4, ax,'line', title_name)

    #ax = plot_stock_data(df_gdp, ax, 'line','2010-2022年国内每季度gdp增速同比')
    print_save_table(df_gdp,'GDP预测',True)

    # show_dynamic_table(dynamic_new)


    # ax = plot_stock_data(res_earning, None, 'bar', '张坤管理各个基金收益率')
    # stock_data = get_index_constituent('上证50', '20230101', '20230508')
    # stock_list = select_value_by_column(stock_data, 'stock_name', -1)
    # res_earning = loop_rank(stock_list, calculate_earning_between_two_time, start, end)
    # res_earning_top_n = rank_index_cross_section(res_earnng, 10, False)
    # ax = plot_stock_data(res_earning_top_n, ax, 'bar', title_name)

    # stock_data = get_index_constituent('上证50', '20230101', '20230508')
    # stock_list = select_value_by_column(stock_data, 'stock_name', -1)
    # res_earning = loop_rank(stock_list, calculate_earning_between_two_time, '20230101', '20230508')
    # res_earning_top_n = rank_index_cross_section(res_earning, 10, False)
    # ax = plot_stock_data(res_earning_top_n, ax, 'bar', title_name)

    # fund_code = query_fund_name_or_code(fund_name, '')
    # fund_data = query_fund_data(fund_code, start, end)
    # fund_index = calculate_stock_index(fund_data, 'adj_nav')
    # ax = plot_stock_data(fund_index, ax, 'line', title_name)
    # fund_df = query_fund_Manager('张坤')
    # fund_code = select_value_by_column(fund_df, 'fund_code', -1)
    # res_earning = loop_rank(fund_code, calculate_earning_between_two_time, start, end, 'adj_nav')
    # ax = plot_stock_data(res_earning, None, 'bar', '张坤管理各个基金收益率')
    # company_df = get_company_info('贵州茅台')
    # print_save_table(company_df,'gzmt', False)




    if ax is not None:
        plt.grid()
        plt.show()



# xxx基金经理管理的几只基金中,收益率最高的那只基金的规模是多少----找基金经理search,按收益率排序rank,找到收益率最高的那个select,显示基金信息 show
# 食品饮料行业中所有股票近十年涨幅最大的股票的信息----找行业search(行业分类--找到行业代码,根据行业代码找到股票成分), 收益率排序rank,找到涨幅最大的那个select,显示股票信息show